import json
from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional
from hf_qwen_client import generate_response
from rag_context import retrieve_document_context
from chat_history import list_conversations, list_messages, format_history

class AnalisisIA(BaseModel):
    redaccion_tecnica: str = Field(..., description="Análisis técnico detallado basado en las micrografías y contexto de la muestra. Máximo 3 párrafos.")
    resumen_ejecutivo: str = Field(..., description="Resumen conciso y al grano para el cliente final sobre los resultados.")
    conformidad_norma: bool = Field(..., description="¿Cumple la muestra con los estándares ASTM E112 en base a los datos proporcionados?")
    observaciones: str = Field(..., description="Observaciones adicionales relevantes.")
    analisis_regiones: Dict[str, str] = Field(
        default_factory=dict,
        description="Breve análisis o comentario de no más de 1 párrafo para cada región identificada en la muestra, usando el nombre de la región como clave."
    )
    conclusion: str = Field(..., description="Conclusión final y de cierre del reporte. Debe ser un párrafo contundente.")

def get_chat_context(user_id: int) -> str:
    try:
        conversations = list_conversations(user_id)
        if not conversations:
            return "No hay conversaciones previas del usuario."
        
        # Tomamos la ultima conversacion del usuario
        latest_conv_id = conversations[0]["id"]
        messages = list_messages(latest_conv_id, user_id, limit=20)
        history = format_history(messages)
        return history
    except Exception as e:
        print(f"Error fetching chat history: {e}")
        return "No se pudo recuperar el historial de chat."

def generate_agentic_report(report_id: int, muestra_id: int, user_id: int, raw_data: dict) -> dict:
    
    chat_context = get_chat_context(user_id)
    
    # RAG: Retrieve ASTM document context
    # Create a search query based on the data
    material = raw_data.get('muestra', {}).get('material', 'metal')
    rag_query = f"norma ASTM E112 tamaño de grano {material}"
    astm_context = retrieve_document_context(rag_query, k=3)
    
    system_prompt = f"""Eres un experto metalurgista especializado en la norma ASTM E112.
Tu tarea es analizar los datos de mediciones de tamaño de grano de una muestra y generar conclusiones técnicas de alto valor.

CONTEXTO NORMATIVO (ASTM):
{astm_context}

CONTEXTO ADICIONAL DEL CLIENTE (Chat reciente):
{chat_context}

Debes proporcionar un JSON válido con el análisis de la muestra que respete estrictamente el esquema solicitado.
Claves requeridas: "redaccion_tecnica", "resumen_ejecutivo", "conformidad_norma" (boolean), "observaciones", "analisis_regiones" (diccionario de region -> texto), "conclusion".
Mantén la redacción técnica precisa pero limitada a 3 párrafos como máximo.
"""

    import copy
    clean_data = copy.deepcopy(raw_data)
    
    # Remove heavy arrays and URLs to avoid exceeding context limits on free API
    if 'material' in clean_data:
        clean_data['material'].pop('muestra_image_url', None)
        clean_data['material'].pop('quality_hist_path', None)
        if 'region_data' in clean_data['material']:
            for region_name, rdata in clean_data['material']['region_data'].items():
                rdata.pop('imagen', None)
                rdata.pop('hist_path', None)
                for micro in rdata.get('micrografias', []):
                    micro.pop('imagen', None)
                    micro.pop('measure_imagen', None)
                    micro.pop('distribution_um', None)

    user_prompt = f"Genera el análisis técnico en JSON puro (sin Markdown) para la siguiente muestra:\n{json.dumps(clean_data, ensure_ascii=False)}"

    try:
        content = generate_response(
            f"{system_prompt}\n\n{user_prompt}",
            max_new_tokens=1500,
        )
        content = content.strip()
        if content.startswith("```json"):
            content = content[7:-3].strip()
        elif content.startswith("```"):
            content = content[3:-3].strip()
            
        analisis = json.loads(content)
        validated_analisis = AnalisisIA(**analisis)
        
        raw_data["analisis_ia"] = validated_analisis.model_dump()
        return raw_data
        
    except Exception as e:
        print(f"Error calling LLM: {e}")
        raw_data["analisis_ia"] = {
            "redaccion_tecnica": "No se pudo generar el análisis automático.",
            "resumen_ejecutivo": "Error en el agente de IA.",
            "conformidad_norma": False,
            "observaciones": str(e),
            "conclusion": "No se pudo generar la conclusión."
        }
        return raw_data