"""Estimador de costos y tokens.

Encapsula la lógica de estimación previa al procesamiento.
"""
import os
from typing import List, Dict, Any, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from application.service import InvoiceProcessingService

from application.profiles import get_profile
from application.prompts import PromptBuilder
from application.normalizer import normalize_to_schema
from application.costs import compute_cost_from_usage, sum_costs, accumulate_usage
from application.file_discovery import get_sorted_radicados, find_file_for_type_in_radicados, get_files_as_radicados
from core.utils import collect_radicado_docs, collect_documents_flat

class Estimator:
    def __init__(self, service: 'InvoiceProcessingService'):
        self.service = service

    def estimate(
        self,
        root_path: str,
        keywords: List[str] | None = None,
        schema: Dict[str, str] | None = None,
        selected_fields: List[str] | None = None,
        *,
        doc_types: Optional[List[Dict[str, str]]] = None,
        schema_by_type: Optional[Dict[str, Dict[str, str]]] = None,
        selected_paths: Optional[List[str]] = None,
        structure_mode: str = "radicados",
    ) -> Dict[str, Any]:
        """Calcula estimado de tokens y costos para toda la carpeta basado en el primer documento."""
        if doc_types:
            return self._estimate_multidoc(root_path, doc_types, schema_by_type, selected_paths, structure_mode)
        return self._estimate_singledoc(root_path, keywords, schema, selected_fields, structure_mode)

    def _estimate_multidoc(self, root_path, doc_types, schema_by_type, selected_paths, structure_mode="radicados"):
        if structure_mode == "general":
            radicados = get_files_as_radicados(root_path)
        else:
            radicados = get_sorted_radicados(root_path)

        if not radicados:
            return {"total_docs": 0, "per_doc": {}, "estimated_total": {}}
        
        prof = get_profile('invoice')
        first_radicado_found = None

        agg_extract_usage = {"prompt_tokens": 0, "completion_tokens": 0, "prompt_tokens_details": {"cached_tokens": 0}}
        agg_normalize_usage = {"prompt_tokens": 0, "completion_tokens": 0, "prompt_tokens_details": {"cached_tokens": 0}}

        for dt in doc_types:
            tname = str(dt.get("name") or "").strip() or "Tipo"
            kw = str(dt.get("keyword") or "").strip()
            
            match, rad_used = find_file_for_type_in_radicados(root_path, radicados, kw)
            if not match: continue
            
            if first_radicado_found is None:
                first_radicado_found = rad_used

            sel_all = selected_paths or []
            selected_paths_for_type = [p.split('/', 1)[1] for p in sel_all if p.startswith(f"{tname}/")]
            sch_t = (schema_by_type or {}).get(tname) or {}
            
            try:
                prompt_t = PromptBuilder.build_for_type(tname, selected_paths_for_type, sch_t, profile_name=prof.get('name', 'invoice'))
            except Exception:
                prompt_t = PromptBuilder.build(selected_paths_for_type, sch_t, profile_name=prof.get('name', 'invoice'))
            
            data, usage_ext = self.service.extract_full_from_pdf(match, instruction=prompt_t)
            self._accumulate_usage(agg_extract_usage, usage_ext)
            
            if isinstance(data, dict): data['radicado'] = rad_used
            else: data = {"radicado": rad_used, "data": data}
            
            norm = normalize_to_schema(data, paths=selected_paths_for_type, profile_name=prof.get('name', 'invoice'))
            if isinstance(norm, dict):
                self._accumulate_usage(agg_normalize_usage, norm.get("_usage"))

        return self._build_result(agg_extract_usage, agg_normalize_usage, len(radicados), first_radicado_found or radicados[0])

    def _estimate_singledoc(self, root_path, keywords, schema, selected_fields, structure_mode="radicados"):
        if structure_mode == "general":
            entries = collect_documents_flat(root_path, keywords)
        else:
            entries = collect_radicado_docs(root_path, keywords or [])
            
        if not entries:
             return {"total_docs": 0, "per_doc": {}, "estimated_total": {}}
        
        first = entries[0]
        prof = get_profile('invoice')
        prompt = PromptBuilder.build(selected_fields or [], schema or {})

        data, usage_ext = self.service.extract_full_from_pdf(first["path"], instruction=prompt)
        if isinstance(data, dict): data['radicado'] = first['radicado']
        else: data = {"radicado": first['radicado'], "data": data}

        norm = normalize_to_schema(data, paths=selected_fields, profile_name=prof.get('name', 'invoice'))
        usage_norm = None
        if isinstance(norm, dict):
            usage_norm = norm.get("_usage")

        return self._build_result(usage_ext, usage_norm, len(entries), first["radicado"])

    def _accumulate_usage(self, agg, new_usage):
        accumulate_usage(agg, new_usage)

    def _build_result(self, usage_ext, usage_norm, total_docs, first_radicado):
        model_ext = getattr(self.service.extractor, "model", None)
        cost_ext = compute_cost_from_usage(usage_ext, "extract", model=model_ext)
        # El normalizador está hardcodeado a gpt-4.1-mini en normalizer.py
        cost_norm = compute_cost_from_usage(usage_norm, "normalize", model="gpt-4.1-mini")
        per_doc_cost = sum_costs(cost_ext, cost_norm)
        estimated_total = {k: round(v * total_docs, 6) for k, v in per_doc_cost.items()}

        return {
            "total_docs": total_docs,
            "first_radicado": first_radicado,
            "extractor_usage": usage_ext,
            "normalizer_usage": usage_norm,
            "per_doc_cost": per_doc_cost,
            "estimated_total_cost": estimated_total,
        }