"""
Pipeline d'extraction de données structurées
"""
import os
import json
import requests
import logging
from typing import Dict, Any

logger = logging.getLogger(__name__)

# Configuration Ollama
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://ollama:11434")
OLLAMA_MODEL = "llama3:8b"

def run(doc_id: str, ctx: dict):
    """
    Extraction de données structurées d'un document
    """
    logger.info(f"Extraction du document {doc_id}")

    try:
        # Récupération des données nécessaires
        extracted_text = ctx.get("extracted_text", "")
        classification = ctx.get("classification", {})
        document_type = classification.get("label", "document_inconnu")

        if not extracted_text:
            raise ValueError("Aucun texte extrait disponible pour l'extraction")

        # Limitation de la taille du texte
        text_sample = extracted_text[:20000]  # Limite plus élevée pour l'extraction

        # Extraction selon le type de document
        extracted_data = _extract_with_ollama(text_sample, document_type)

        # Validation des données extraites
        validated_data = _validate_extracted_data(extracted_data, document_type)

        # Stockage du résultat
        ctx["extracted_data"] = validated_data

        # Métadonnées d'extraction
        extract_meta = {
            "extraction_completed": True,
            "document_type": document_type,
            "fields_extracted": len(validated_data),
            "model_used": OLLAMA_MODEL
        }

        ctx["extract_meta"] = extract_meta

        logger.info(f"Extraction terminée pour le document {doc_id}: {len(validated_data)} champs extraits")

    except Exception as e:
        logger.error(f"Erreur lors de l'extraction du document {doc_id}: {e}")
        raise

def _extract_with_ollama(text: str, document_type: str) -> Dict[str, Any]:
    """
    Extraction de données avec Ollama selon le type de document
    """
    try:
        # Chargement du prompt d'extraction
        prompt = _load_extraction_prompt(document_type)

        # Remplacement du placeholder
        full_prompt = prompt.replace("{{TEXT}}", text)

        # Appel à l'API Ollama
        payload = {
            "model": OLLAMA_MODEL,
            "prompt": full_prompt,
            "stream": False,
            "options": {
                "temperature": 0.1,
                "top_p": 0.9,
                "max_tokens": 1000
            }
        }

        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/generate",
            json=payload,
            timeout=180
        )

        if response.status_code != 200:
            raise RuntimeError(f"Erreur API Ollama: {response.status_code} - {response.text}")

        result = response.json()

        # Parsing de la réponse JSON
        try:
            extracted_data = json.loads(result["response"])
        except json.JSONDecodeError:
            # Fallback si la réponse n'est pas du JSON valide
            extracted_data = _parse_fallback_extraction(result["response"], document_type)

        return extracted_data

    except Exception as e:
        logger.error(f"Erreur lors de l'extraction avec Ollama: {e}")
        return {"error": str(e), "extraction_failed": True}

def _load_extraction_prompt(document_type: str) -> str:
    """
    Chargement du prompt d'extraction selon le type de document
    """
    prompt_path = f"/app/models/prompts/extract_{document_type}_prompt.txt"

    try:
        if os.path.exists(prompt_path):
            with open(prompt_path, 'r', encoding='utf-8') as f:
                return f.read()
    except Exception as e:
        logger.warning(f"Impossible de charger le prompt d'extraction pour {document_type}: {e}")

    # Prompt générique par défaut
    return _get_generic_extraction_prompt()

def _get_generic_extraction_prompt() -> str:
    """
    Prompt générique d'extraction
    """
    return """
Tu es un expert en extraction de données notariales. Analyse le texte suivant et extrais les informations importantes.

TEXTE À ANALYSER :
{{TEXT}}

Extrais les informations suivantes si elles sont présentes :
- dates importantes
- montants financiers
- noms de personnes
- adresses
- références de biens
- numéros de documents

Réponds UNIQUEMENT avec un JSON valide :
{
  "dates": ["date1", "date2"],
  "montants": ["montant1", "montant2"],
  "personnes": ["nom1", "nom2"],
  "adresses": ["adresse1", "adresse2"],
  "references": ["ref1", "ref2"],
  "notes": "informations complémentaires"
}
"""

def _validate_extracted_data(data: Dict[str, Any], document_type: str) -> Dict[str, Any]:
    """
    Validation des données extraites
    """
    if not isinstance(data, dict):
        return {"error": "Données extraites invalides", "raw_data": str(data)}

    # Validation selon le type de document
    if document_type == "acte_vente":
        return _validate_vente_data(data)
    elif document_type == "acte_achat":
        return _validate_achat_data(data)
    elif document_type == "donation":
        return _validate_donation_data(data)
    elif document_type == "testament":
        return _validate_testament_data(data)
    elif document_type == "succession":
        return _validate_succession_data(data)
    else:
        return _validate_generic_data(data)

def _validate_vente_data(data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Validation des données d'acte de vente
    """
    validated = {
        "type": "acte_vente",
        "vendeur": data.get("vendeur", ""),
        "acheteur": data.get("acheteur", ""),
        "bien": data.get("bien", ""),
        "prix": data.get("prix", ""),
        "date_vente": data.get("date_vente", ""),
        "notaire": data.get("notaire", ""),
        "etude": data.get("etude", ""),
        "adresse_bien": data.get("adresse_bien", ""),
        "surface": data.get("surface", ""),
        "references": data.get("references", []),
        "notes": data.get("notes", "")
    }

    return validated

def _validate_achat_data(data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Validation des données d'acte d'achat
    """
    validated = {
        "type": "acte_achat",
        "vendeur": data.get("vendeur", ""),
        "acheteur": data.get("acheteur", ""),
        "bien": data.get("bien", ""),
        "prix": data.get("prix", ""),
        "date_achat": data.get("date_achat", ""),
        "notaire": data.get("notaire", ""),
        "etude": data.get("etude", ""),
        "adresse_bien": data.get("adresse_bien", ""),
        "surface": data.get("surface", ""),
        "references": data.get("references", []),
        "notes": data.get("notes", "")
    }

    return validated

def _validate_donation_data(data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Validation des données de donation
    """
    validated = {
        "type": "donation",
        "donateur": data.get("donateur", ""),
        "donataire": data.get("donataire", ""),
        "bien_donne": data.get("bien_donne", ""),
        "valeur": data.get("valeur", ""),
        "date_donation": data.get("date_donation", ""),
        "notaire": data.get("notaire", ""),
        "etude": data.get("etude", ""),
        "conditions": data.get("conditions", ""),
        "references": data.get("references", []),
        "notes": data.get("notes", "")
    }

    return validated

def _validate_testament_data(data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Validation des données de testament
    """
    validated = {
        "type": "testament",
        "testateur": data.get("testateur", ""),
        "heritiers": data.get("heritiers", []),
        "legs": data.get("legs", []),
        "date_testament": data.get("date_testament", ""),
        "notaire": data.get("notaire", ""),
        "etude": data.get("etude", ""),
        "executeur": data.get("executeur", ""),
        "references": data.get("references", []),
        "notes": data.get("notes", "")
    }

    return validated

def _validate_succession_data(data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Validation des données de succession
    """
    validated = {
        "type": "succession",
        "defunt": data.get("defunt", ""),
        "heritiers": data.get("heritiers", []),
        "biens": data.get("biens", []),
        "date_deces": data.get("date_deces", ""),
        "date_partage": data.get("date_partage", ""),
        "notaire": data.get("notaire", ""),
        "etude": data.get("etude", ""),
        "references": data.get("references", []),
        "notes": data.get("notes", "")
    }

    return validated

def _validate_generic_data(data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Validation générique des données
    """
    validated = {
        "type": "document_generique",
        "dates": data.get("dates", []),
        "montants": data.get("montants", []),
        "personnes": data.get("personnes", []),
        "adresses": data.get("adresses", []),
        "references": data.get("references", []),
        "notes": data.get("notes", "")
    }

    return validated

def _parse_fallback_extraction(response_text: str, document_type: str) -> Dict[str, Any]:
    """
    Parsing de fallback pour l'extraction
    """
    # Extraction basique avec regex
    import re

    # Extraction des dates
    dates = re.findall(r'\b\d{1,2}[/-]\d{1,2}[/-]\d{2,4}\b', response_text)

    # Extraction des montants
    amounts = re.findall(r'\b\d{1,3}(?:\s\d{3})*(?:[.,]\d{2})?\s*€\b', response_text)

    # Extraction des noms (basique)
    names = re.findall(r'\b(?:Monsieur|Madame|M\.|Mme\.)\s+[A-Z][a-z]+', response_text)

    return {
        "dates": dates,
        "montants": amounts,
        "personnes": names,
        "extraction_method": "fallback",
        "document_type": document_type
    }