#!/usr/bin/env python3
"""
OCR Pipeline para Documentos Contabilísticos
PDF → imagem → RapidOCR → DeepSeek → JSON estruturado

Uso: python3 ocr-invoice.py <caminho_do_pdf>
"""
import sys
import os
import json
import time
from pathlib import Path

# ── CONFIG ───────────────────────────────────────────────────────
OPENCODE_GO_KEY = ""
# Ler .env automaticamente
for line in open(os.path.expanduser("~/.hermes/.env")):
    line = line.strip()
    if line.startswith("OPENCODE_GO_API_KEY=") and not line.startswith("#"):
        OPENCODE_GO_KEY = line.split("=", 1)[1]
        break
OPENCODE_GO_URL = "https://opencode.ai/zen/go/v1"
MODEL = "deepseek-v4-flash"
OCR_DPI = 200


def pdf_to_image(pdf_path: str, dpi: int = OCR_DPI) -> str:
    """Converte primeira página do PDF para imagem PNG."""
    import pymupdf
    doc = pymupdf.open(pdf_path)
    page = doc[0]
    pix = page.get_pixmap(dpi=dpi)
    img_path = f"/tmp/ocr_{Path(pdf_path).stem}.png"
    pix.save(img_path)
    return img_path


def ocr_image(img_path: str) -> list[dict]:
    """Executa RapidOCR na imagem. Retorna lista de {text, score, bbox}."""
    from rapidocr import RapidOCR
    engine = RapidOCR()
    result = engine(img_path)
    
    lines = []
    if result and result.txts:
        for txt, score in zip(result.txts, result.scores):
            lines.append({"text": txt, "score": float(score)})
    return lines


def extract_structured(ocr_lines: list[dict], pdf_name: str) -> dict:
    """Usa DeepSeek para extrair campos estruturados do texto OCR."""
    from openai import OpenAI
    
    client = OpenAI(
        api_key=OPENCODE_GO_KEY,
        base_url=OPENCODE_GO_URL,
    )
    
    # Texto OCR como bloco
    ocr_text = "\n".join(f"[{l['score']:.2f}] {l['text']}" for l in ocr_lines)
    
    prompt = f"""Analisa o seguinte texto extraído de um documento contabilístico (factura/recibo) via OCR.
Extrai os campos estruturados e devolve APENAS JSON válido (sem markdown, sem ```).

Texto OCR:
{ocr_text}

Nome do ficheiro: {pdf_name}

Devolve JSON com esta estrutura exata:
{{
  "tipo_documento": "factura|recibo|nota_de_credito|outro",
  "fornecedor": {{
    "nome": "string",
    "nif_cif": "string",
    "morada": "string",
    "telefone": "string"
  }},
  "cliente": {{
    "nome": "string",
    "nif": "string",
    "morada": "string",
    "telefone": "string"
  }},
  "documento": {{
    "numero": "string",
    "data": "YYYY-MM-DD",
    "metodo_pagamento": "string",
    "referencia": "string"
  }},
  "artigos": [
    {{
      "codigo": "string",
      "descricao": "string",
      "preco_unitario": 0.00,
      "quantidade": 1,
      "total": 0.00,
      "notas": "string"
    }}
  ],
  "resumo": {{
    "base_tributavel": 0.00,
    "taxa_iva_percent": 0,
    "iva_valor": 0.00,
    "recargo_percent": 0,
    "recargo_valor": 0.00,
    "total_pagar": 0.00,
    "moeda": "EUR"
  }},
  "notas": "string com observações relevantes"
}}

Regras:
- Preços com vírgula decimal (formato PT: 1.234,56 → 1234.56)
- Se campo não encontrado, usar null
- Se artigos não detectados, array vazio
- IVA: se não explícito, calcular a partir de base + total
- Moeda: EUR por defeito"""

    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": "És um assistente especializado em extração de dados de documentos contabilísticos portugueses. Devolves sempre JSON válido."},
            {"role": "user", "content": prompt}
        ],
        temperature=0.0,
        max_tokens=16384,
    )
    
    raw = response.choices[0].message.content.strip()
    
    # Limpar possíveis wrappers markdown
    if raw.startswith("```"):
        raw = raw.split("\n", 1)[1]
        if raw.endswith("```"):
            raw = raw[:-3]
        raw = raw.strip()
    
    return json.loads(raw)


def process_invoice(pdf_path: str) -> dict:
    """Pipeline completo: PDF → JSON estruturado."""
    print(f"📄 A processar: {pdf_path}")
    
    t0 = time.time()
    
    # 1. PDF → imagem
    print("  [1/3] PDF → imagem...")
    img = pdf_to_image(pdf_path)
    print(f"        OK ({time.time()-t0:.1f}s)")
    
    # 2. OCR
    print("  [2/3] OCR (RapidOCR)...")
    t1 = time.time()
    lines = ocr_image(img)
    print(f"        {len(lines)} linhas em {time.time()-t1:.1f}s")
    
    # 3. Structured extraction
    print("  [3/3] Extração estruturada (DeepSeek)...")
    t2 = time.time()
    structured = extract_structured(lines, os.path.basename(pdf_path))
    print(f"        OK ({time.time()-t2:.1f}s)")
    
    elapsed = time.time() - t0
    print(f"\n✅ Pipeline completo em {elapsed:.1f}s")
    
    # Cleanup
    os.remove(img)
    
    return {
        "source_file": pdf_path,
        "ocr_lines": len(lines),
        "processing_time_seconds": round(elapsed, 1),
        "extracted_data": structured,
    }


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Uso: python3 {sys.argv[0]} <caminho_do_pdf>")
        sys.exit(1)
    
    pdf_path = sys.argv[1]
    if not os.path.exists(pdf_path):
        print(f"Erro: ficheiro não encontrado: {pdf_path}")
        sys.exit(1)
    
    result = process_invoice(pdf_path)
    print("\n" + json.dumps(result, indent=2, ensure_ascii=False))