"""
Extractor de Edictos Societarios — Constitución y Cambios de Directorio
========================================================================
Requisitos:
    pip install google-genai python-dotenv pypdf

Configuración:
    .env con GEMINI_API_KEY=...

Uso:
    python3 sociedades.py boletin.pdf
    python3 sociedades.py boletin.pdf --model gemini-2.5-flash-lite
    python3 sociedades.py boletin.pdf --estimate
"""

import logging
import sys
import os
import re
import json
import time
import pathlib
import unicodedata
from datetime import datetime
from dotenv import load_dotenv
from google import genai
from google.genai import types
from google.genai.errors import ClientError, ServerError

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent.parent))
from helpers.logging_setup import setup_logging

load_dotenv()
setup_logging("vertex")

# ── Configuración ──────────────────────────────────────────────
MODELS_FALLBACK = [
    "gemini-3-flash-preview",
    "gemini-3.1-flash-lite-preview",
    "gemini-2.5-flash"
    #"gemini-2.5-flash-lite",
    #"gemini-2.5-pro",
]
MAX_RETRIES = 5
TEXT_BATCH_SIZE = 25  # edictos por batch en el pass 2

# ── Contexto de dominio ────────────────────────────────────────
CONTEXTO = """
En Argentina, el Boletín Oficial es la publicación oficial donde se dan a publicidad actos jurídicos
que requieren notificación pública para tener validez frente a terceros.

Un EDICTO DE CONSTITUCIÓN se publica cuando se crea una sociedad comercial (S.A., S.R.L., etc.),
exigido por la Ley General de Sociedades 19.550. Contiene: denominación y tipo societario, domicilio,
objeto social, capital social, datos de socios/accionistas fundadores, órganos de administración y
fiscalización, y fecha del estatuto constitutivo.

Un EDICTO DE CAMBIO DE DIRECTORIO se publica cuando una sociedad modifica su órgano de administración
(Directorio en S.A., Gerencia en S.R.L.) por designaciones, renuncias, remociones o renovaciones.
Contiene: nombre de la sociedad, directores/gerentes salientes y entrantes, cargos, duración del
mandato y fecha de la asamblea que resolvió el cambio.

Un EDICTO DE DISOLUCIÓN se publica cuando una sociedad decide disolverse y entrar en proceso de
liquidación voluntaria. Contiene: nombre de la sociedad, fecha de la asamblea que resolvió la
disolución, datos del liquidador designado y domicilio especial.

Un EDICTO DE CESIÓN DE CUOTAS se publica cuando uno o más socios de una S.R.L. transfieren su
participación a terceros o a otros socios. Contiene: nombre de la sociedad, cedente/s, cesionario/s,
cantidad de cuotas cedidas y precio o condiciones de la cesión.

Un EDICTO DE LLAMADO A ASAMBLEA se publica para convocar a los socios o accionistas a una asamblea
ordinaria o extraordinaria. Contiene: nombre de la sociedad, tipo de asamblea, fecha, hora y lugar
de celebración, y orden del día.

Un EDICTO DE AUMENTO DE CAPITAL se publica cuando una sociedad resuelve incrementar su capital social.
Contiene: nombre de la sociedad, monto del capital anterior, monto del nuevo capital, tipo de aumento
(nuevos aportes, capitalización de reservas, etc.) y fecha de la asamblea que lo resolvió.
"""

# ── Prompt Pass 1: detectar edictos sin texto ──────────────────
PROMPT_META = f"""
{CONTEXTO}

Analizá este boletín oficial y realizá dos tareas:

1. Encontrá la fecha de publicación del boletín (suele estar en el encabezado o tapa).
   Devolvela en formato ISO: YYYY-MM-DD. Si no la encontrás, usá null.

2. Encontrá ÚNICAMENTE los edictos de estos tipos:
   - CONSTITUCIÓN DE SOCIEDAD: cuando se crea una nueva S.A., S.R.L. u otro tipo societario.
   - CAMBIO DE DIRECTORIO / AUTORIDADES: cuando una sociedad designa, renueva o remueve directores, gerentes u otras autoridades.
   - DISOLUCIÓN: cuando una sociedad resuelve disolverse y entrar en liquidación voluntaria.
   - CESIÓN DE CUOTAS: cuando socios de una S.R.L. transfieren su participación.
   - LLAMADO A ASAMBLEA: convocatoria a asamblea ordinaria o extraordinaria.
   - AUMENTO DE CAPITAL: cuando una sociedad incrementa su capital social.
   Ignorá completamente cualquier otro tipo de edicto.

Para cada edicto encontrado, extraé:
- "razon_social": nombre completo de la sociedad tal como aparece en el PDF
- "tipo": exactamente uno de los valores del enum (ver schema)
- "texto_inicio": las primeras 80 palabras del texto del edicto, para poder identificarlo después
- "provincia": provincia argentina a la que pertenece el edicto, en mayúsculas y sin acentos
  (ej: "BUENOS AIRES", "CORDOBA", "SANTA FE", "ENTRE RIOS", "RIO NEGRO", "SANTIAGO DEL ESTERO",
  "TIERRA DEL FUEGO", "SAN LUIS", "SAN JUAN", "LA RIOJA", "LA PAMPA", "CIUDAD AUTONOMA DE BUENOS AIRES").
  Si no podés determinarla con certeza, devolvé null.
"""

SCHEMA_META = {
    "type": "OBJECT",
    "properties": {
        "fecha_publicacion": {
            "type": "STRING",
            "nullable": True,
            "description": "Fecha de publicación del boletín en formato YYYY-MM-DD, null si no se encuentra"
        },
        "edictos": {
            "type": "ARRAY",
            "items": {
                "type": "OBJECT",
                "properties": {
                    "razon_social": {
                        "type": "STRING",
                        "description": "Nombre completo de la sociedad"
                    },
                    "tipo": {
                        "type": "STRING",
                        "enum": ["constitución", "cambio_directorio", "disolución", "cesión_cuotas", "llamado_asamblea", "aumento_capital"],
                        "description": "Tipo de edicto societario"
                    },
                    "texto_inicio": {
                        "type": "STRING",
                        "description": "Primeras 80 palabras del texto del edicto"
                    },
                    "provincia": {
                        "type": "STRING",
                        "nullable": True,
                        "description": "Provincia argentina en mayúsculas y sin acentos. Null si no se puede determinar."
                    },
                },
                "required": ["razon_social", "tipo", "texto_inicio", "provincia"]
            }
        }
    },
    "required": ["fecha_publicacion", "edictos"]
}

# ── Prompt Pass 2: obtener texto de edictos específicos ────────
def build_prompt_texto(edictos_meta: list[dict]) -> str:
    items = "\n".join(
        f"- {e['razon_social']} | inicio: \"{e.get('texto_inicio', '')[:120]}\""
        for e in edictos_meta
    )
    return f"""
{CONTEXTO}

Del boletín oficial, extraé el texto COMPLETO e ÍNTEGRO de los siguientes edictos societarios.
Para identificar cada edicto usá tanto la razón social como el fragmento de texto de inicio:

{items}

Para cada uno devolvé:
- "razon_social": el nombre exactamente como figura en el PDF
- "texto_edicto": el texto íntegro del edicto, palabra por palabra, sin resumir ni recortar

Si no encontrás alguno, omitilo del resultado.
"""

SCHEMA_TEXTO = {
    "type": "ARRAY",
    "items": {
        "type": "OBJECT",
        "properties": {
            "razon_social": {
                "type": "STRING",
                "description": "Nombre de la sociedad tal como figura en el PDF"
            },
            "texto_edicto": {
                "type": "STRING",
                "description": "Texto completo e íntegro del edicto"
            },
        },
        "required": ["razon_social", "texto_edicto"]
    }
}

# ── Precios ────────────────────────────────────────────────────
MODEL_PRICING = {
    "gemini-3.1-flash-lite-preview":     (0.25,  1.50),
    "gemini-3-flash-preview":            (0.50,  3.00),
    "gemini-3.1-pro-preview":            (2.00, 12.00),
    "gemini-2.5-flash":                  (0.30,  2.50),
    "gemini-2.5-flash-lite":             (0.10,  0.40),
    "gemini-2.5-pro":                    (1.25, 10.00),
    "gemini-2.0-flash":                  (0.10,  0.40),
}


# ── Helpers ────────────────────────────────────────────────────
class ModelUnavailableError(Exception):
    pass


def extract_retry_delay(error_message: str) -> float:
    match = re.search(r"retry in ([\d.]+)s", str(error_message), re.IGNORECASE)
    return float(match.group(1)) if match else 30.0


def is_quota_exhausted_daily(error_message: str) -> bool:
    return "limit: 0" in str(error_message)


def call_api(client, models, contents, config):
    """Llama a la API con retry y fallback entre modelos."""
    last_error = None
    for model in models:
        for attempt in range(1, MAX_RETRIES + 1):
            try:
                chunks = []
                for chunk in client.models.generate_content_stream(
                    model=model, contents=contents, config=config
                ):
                    if chunk.text:
                        chunks.append(chunk.text)
                raw = "".join(chunks).strip()
                raw = re.sub(r"^```json\s*", "", raw)
                raw = re.sub(r"\s*```$", "", raw)
                return json.loads(raw), model
            except ClientError as e:
                if e.code == 429:
                    if is_quota_exhausted_daily(str(e)):
                        logging.warning(f"CUOTA_DIARIA_AGOTADA | modelo={model}")
                        break
                    delay = extract_retry_delay(str(e))
                    logging.warning(f"REINTENTO | rate_limit intento={attempt}/{MAX_RETRIES} delay={delay:.0f}s")
                    time.sleep(delay)
                else:
                    raise
            except ServerError:
                delay = 15.0 * attempt
                logging.warning(f"REINTENTO | servidor_no_disponible intento={attempt}/{MAX_RETRIES} delay={delay:.0f}s")
                time.sleep(delay)
            except json.JSONDecodeError as e:
                logging.warning(f"JSON_INVALIDO | {e}")
                last_error = e
                break
        last_error = last_error or Exception(f"Reintentos agotados para {model}")
        if len(models) > 1:
            logging.info("MODELO_FALLBACK | probando siguiente modelo")
    raise Exception(f"No se pudo completar la llamada. Último error: {last_error}")


def make_contents(pdf_bytes, prompt):
    return [
        types.Part.from_bytes(data=pdf_bytes, mime_type="application/pdf"),
        prompt,
    ]


def make_config(schema):
    return types.GenerateContentConfig(
        response_mime_type="application/json",
        response_schema=schema,
        max_output_tokens=65536,
    )


# ── Lógica two-pass ────────────────────────────────────────────
def extract_sociedades(client, models, pdf_path):
    pdf_bytes = pdf_path.read_bytes()

    # ── Pass 1: detectar edictos + fecha de publicación ──
    logging.info("PASS_1 | detectando edictos societarios y fecha del boletín")
    pass1, model_used = call_api(
        client, models,
        make_contents(pdf_bytes, PROMPT_META),
        make_config(SCHEMA_META),
    )
    fecha_publicacion = pass1.get("fecha_publicacion")
    meta_list = pass1.get("edictos", [])
    logging.info(f"PASS_1_OK | fecha={fecha_publicacion or 'no encontrada'} edictos={len(meta_list)} modelo={model_used}")

    if not meta_list:
        return [], fecha_publicacion, model_used

    # ── Pass 2: obtener texto en batches ──
    logging.info(f"PASS_2 | obteniendo textos en batches de {TEXT_BATCH_SIZE}")
    textos_map = {}
    batches = [meta_list[i:i+TEXT_BATCH_SIZE] for i in range(0, len(meta_list), TEXT_BATCH_SIZE)]

    for i, batch in enumerate(batches, 1):
        logging.info(f"PASS_2_BATCH | {i}/{len(batches)} ({len(batch)} edictos)")
        texto_list, model_used = call_api(
            client, models,
            make_contents(pdf_bytes, build_prompt_texto(batch)),
            make_config(SCHEMA_TEXTO),
        )
        logging.info(f"PASS_2_BATCH_OK | {len(texto_list)} textos modelo={model_used}")
        for item in texto_list:
            textos_map[item["razon_social"].strip()] = item.get("texto_edicto", "")

    # ── Merge pass 1 + pass 2 (matching flexible) ──
    def normalizar(s: str) -> str:
        """Quita tildes, puntos, espacios extra y pasa a mayúsculas para comparar."""
        s = unicodedata.normalize("NFKD", s)
        s = "".join(c for c in s if not unicodedata.combining(c))
        s = s.upper()
        s = re.sub(r'[.\-]', '', s)   # quita puntos y guiones (S.R.L. → SRL)
        s = re.sub(r'\s+', ' ', s).strip()
        return s

    # Reconstruir textos_map con claves normalizadas
    textos_map_norm = {normalizar(k): v for k, v in textos_map.items()}

    def find_texto(razon_social: str) -> str:
        key = normalizar(razon_social)
        # Coincidencia exacta normalizada
        if key in textos_map_norm:
            return textos_map_norm[key]
        # Coincidencia parcial normalizada
        for k, v in textos_map_norm.items():
            if key in k or k in key:
                return v
        return ""

    resultado = []
    for meta in meta_list:
        resultado.append({
            "razon_social": meta["razon_social"],
            "tipo": meta["tipo"],
            "provincia": meta.get("provincia"),
            "texto_edicto": find_texto(meta["razon_social"]),
        })

    return resultado, fecha_publicacion, model_used


# ── Estimación de costo ────────────────────────────────────────
def estimate_cost(client, models, pdf_path):
    model = models[0]
    price_in, price_out = MODEL_PRICING.get(model, (0.30, 2.50))
    pdf_bytes = pdf_path.read_bytes()

    print(f"Modelo: {model}  |  ${price_in}/1M in · ${price_out}/1M out")
    print(f"{'─' * 55}")

    # Tokens pass 1
    contents_meta = make_contents(pdf_bytes, PROMPT_META)
    resp1 = client.models.count_tokens(model=model, contents=contents_meta)
    tok_pass1 = resp1.total_tokens
    print(f"  Pass 1 (metadata, PDF completo):   {tok_pass1:>10,} tokens")

    # Tokens pass 2 (estimado: asumimos 50 edictos, 1 batch de ejemplo)
    ejemplo_batch = [
        {"razon_social": "EJEMPLO S.A.", "texto_inicio": "Se constituye la sociedad..."},
        {"razon_social": "OTRA SOCIEDAD S.R.L.", "texto_inicio": "Por escritura pública..."},
    ]
    contents_texto = make_contents(pdf_bytes, build_prompt_texto(ejemplo_batch))
    resp2 = client.models.count_tokens(model=model, contents=contents_texto)
    tok_por_batch = resp2.total_tokens
    batches_estimados = 3  # estimado conservador
    tok_pass2_total = tok_por_batch * batches_estimados
    print(f"  Pass 2 (~{batches_estimados} batches, {TEXT_BATCH_SIZE} edictos c/u): {tok_pass2_total:>10,} tokens est.")

    total_input = tok_pass1 + tok_pass2_total
    total_output_est = 65536  # máximo conservador
    cost_in  = total_input      / 1_000_000 * price_in
    cost_out = total_output_est / 1_000_000 * price_out
    print(f"{'─' * 55}")
    print(f"  Total tokens entrada (est.):       {total_input:>10,}")
    print(f"  Costo entrada:                     ${cost_in:.6f}")
    print(f"  Costo output (est. máximo):        ${cost_out:.6f}")
    print(f"  COSTO TOTAL ESTIMADO:              ${cost_in+cost_out:.6f}")


# ── Output ─────────────────────────────────────────────────────
def save_json(edictos, fecha_publicacion, pdf_name, output_dir=None):
    stem = pathlib.Path(pdf_name).stem
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"sociedades_{stem}_{timestamp}.json"
    base = pathlib.Path(output_dir) if output_dir else pathlib.Path(".")
    base.mkdir(parents=True, exist_ok=True)
    filepath = base / filename
    output = {
        "pdf_origen": pdf_name,
        "fecha_publicacion": fecha_publicacion,
        "fecha_procesamiento": datetime.now().isoformat(),
        "total": len(edictos),
        "edictos": edictos,
    }
    filepath.write_text(json.dumps(output, indent=2, ensure_ascii=False), encoding="utf-8")
    return filepath


def print_results(edictos, fecha_publicacion, pdf_name):
    if not edictos:
        print("\nNo se encontraron edictos societarios.")
        return

    constituciones = [e for e in edictos if e["tipo"] == "constitución"]
    cambios        = [e for e in edictos if e["tipo"] == "cambio_directorio"]

    print(f"\n{'=' * 90}")
    print(f"  EDICTOS SOCIETARIOS: {len(edictos)}  "
          f"(constituciones: {len(constituciones)}  |  cambios de directorio: {len(cambios)})")
    print(f"  PDF: {pdf_name}")
    print(f"  Fecha publicación: {fecha_publicacion or 'no encontrada'}")
    print(f"{'=' * 90}\n")

    print(f"{'#':<4} {'RAZÓN SOCIAL':<45} {'TIPO'}")
    print(f"{'-'*4} {'-'*45} {'-'*20}")
    for i, e in enumerate(edictos, 1):
        razon = (e.get("razon_social") or "N/A")[:44]
        tipo  = e.get("tipo", "N/A")
        print(f"{i:<4} {razon:<45} {tipo}")

    print(f"\n{'=' * 90}")
    print("  DETALLE")
    print(f"{'=' * 90}")
    for i, e in enumerate(edictos, 1):
        texto = e.get("texto_edicto", "")
        preview = texto[:400] + "..." if len(texto) > 400 else texto
        print(f"\n  [{i}] {e.get('razon_social', 'N/A')}  ({e.get('tipo', '')})")
        print(f"      Texto: {preview}")


# ── Main ───────────────────────────────────────────────────────
def main():
    if len(sys.argv) < 2:
        print("Uso: python3 sociedades.py <boletin.pdf|directorio> [--model modelo] [--estimate] [--out directorio]")
        sys.exit(1)

    input_path = pathlib.Path(sys.argv[1])
    if not input_path.exists():
        print(f"No se encontró: {input_path}")
        sys.exit(1)

    explicit_model = None
    estimate = False
    output_dir = str(pathlib.Path(__file__).resolve().parent.parent.parent / "files" / "boletines_json")
    args = sys.argv[2:]

    if "--model" in args:
        idx = args.index("--model")
        explicit_model = args[idx + 1]
        args = args[:idx] + args[idx + 2:]

    if "--estimate" in args:
        estimate = True
        args.remove("--estimate")

    if "--out" in args:
        idx = args.index("--out")
        output_dir = args[idx + 1]
        args = args[:idx] + args[idx + 2:]

    api_key = os.getenv("GEMINI_API_KEY", "")
    if not api_key:
        print("Falta GEMINI_API_KEY en .env")
        sys.exit(1)

    if explicit_model:
        models = [explicit_model] + [m for m in MODELS_FALLBACK if m != explicit_model]
    else:
        models = MODELS_FALLBACK

    # Armar lista de PDFs a procesar
    if input_path.is_dir():
        pdfs = sorted(input_path.glob("*.pdf"))
        if not pdfs:
            print(f"No se encontraron archivos PDF en: {input_path}")
            sys.exit(1)
        # Si no se especificó --out, guardar los JSON en la misma carpeta
        if output_dir is None:
            output_dir = str(input_path)
    else:
        pdfs = [input_path]

    client = genai.Client(api_key=api_key)
    logging.info(f"PROCESO_INICIADO | modelos={' → '.join(models)} pdfs={len(pdfs)} output={output_dir}")

    for i, pdf_path in enumerate(pdfs, 1):
        size_mb = pdf_path.stat().st_size / (1024 * 1024)
        logging.info(f"PDF_INICIO | [{i}/{len(pdfs)}] {pdf_path.name} size={size_mb:.1f}MB")

        if estimate:
            estimate_cost(client, models, pdf_path)
            continue

        try:
            edictos, fecha_publicacion, _ = extract_sociedades(client, models, pdf_path)
            print_results(edictos, fecha_publicacion, pdf_path.name)
            json_path = save_json(edictos, fecha_publicacion, pdf_path.name, output_dir)
            procesados_dir = pdf_path.parent / "procesados"
            procesados_dir.mkdir(exist_ok=True)
            pdf_path.rename(procesados_dir / pdf_path.name)
            logging.info(f"PDF_OK | {pdf_path.name} edictos={len(edictos)} fecha={fecha_publicacion} json={json_path}")
        except Exception as e:
            logging.error(f"PDF_ERROR | {pdf_path.name}: {e}")


if __name__ == "__main__":
    main()