"""
Scenario Generator — Talks to Gemini using LangChain and strict Pydantic parsing.
"""

import json
import uuid
from typing import Optional
from langchain_groq import ChatGroq
from config import get_settings
from models import ZScoreResult, ScenarioGenerationResult

SYSTEM_PROMPT = """You are a Financial Scenario Generator AI.
Your task is to create high-quality, pedagogical financial scenarios using historical market data and Z-score analysis.

GENERATE EXACTLY 5 DIVERSE SCENARIOS from the provided events. 

You MUST return ONLY valid JSON matching this exact structure:
{
  "totalPossibleScenarios": 5,
  "scenarios": [
    {
      "id": "10-character string (e.g. SCEN-00A3X)",
      "title": "Clear, descriptive title",
      "shortDescription": "1-2 sentence overview",
      "givensTable": {
        "date": "...",
        "stockSymbol": "...",
        "price": 0.0,
        "zScore": 0.0,
        "marketConditions": "...",
        "eventType": "...",
        "context": "..."
      },
      "scenarioParagraph": "Detailed narrative describing the market situation... Clearly mention whether this is a MAJOR CRISIS EVENT or NORMAL MARKET VOLATILITY.",
      "bestAnswer": {
        "answer": "...",
        "rationale": "..."
      },
      "otherAnswers": [
        { "answer": "...", "explanation": "..." },
        { "answer": "...", "explanation": "..." },
        { "answer": "...", "explanation": "..." }
      ],
      "riskLevel": "Low | Medium | High"
    }
  ]
}

Event-Type Classification Rules:
- "major" → context contains crisis keywords (COVID, Crash, Financial Crisis, pandemic, collapse, 9/11) OR |Z-score| >= 3.0
- "normal" → everything else

Risk Classification Rules:
- "High" → Major market crashes, high volatility during crises.
- "Medium" → Notable daily volatility or uncertainty.
- "Low" → Slight corrections or rallies in otherwise stable periods.

Return ONLY valid JSON. No markdown fences, no commentary, no additional text outside the JSON object.
"""

def _build_user_prompt(
    stock_symbol: str,
    zscore_result: ZScoreResult,
) -> str:
    events_data = [e.model_dump() for e in zscore_result.events]
    return (
        f"Stock Symbol: {stock_symbol}\n"
        f"Total Events Available: {zscore_result.total_events}\n"
        f"Window Size: {zscore_result.window_size}\n"
        f"Data Points: {zscore_result.data_points}\n\n"
        f"Events Data:\n{json.dumps(events_data, indent=2)}"
    )

def generate_scenarios(stock_symbol: str, zscore_result: ZScoreResult, knowledge_base_id: Optional[str] = None) -> ScenarioGenerationResult:
    """
    Call Groq API → parse structured JSON → return validated Pydantic model.
    """
    settings = get_settings()
    
    # Initialize Groq client
    llm = ChatGroq(
        api_key=settings.GROQ_API_KEY,
        model_name=settings.GROQ_MODEL,
        temperature=0.7,
        max_tokens=8000
    )

    user_prompt = _build_user_prompt(stock_symbol, zscore_result)
    
    system_message = SYSTEM_PROMPT
    if knowledge_base_id:
        from services.rag_engine import search_kb
        style_context = search_kb(knowledge_base_id, "financial investment scenario examples style guide writing", top_k=2)
        system_message += f"\n\n--- RAG STYLE EXAMPLES & CONTEXT ---\nPlease mimic the tone, structure, or terminology from this retrieved knowledge base if applicable:\n{style_context}\n------------------------------------"
    
    # Send message to Groq
    messages = [
        ("system", system_message),
        ("human", user_prompt),
    ]
    
    response = llm.invoke(messages)
    raw_text = response.content.strip()

    # Parse & Validate
    try:
        data = json.loads(raw_text)
    except json.JSONDecodeError as e:
        # Fallback if markdown fences sneaked in
        import re
        match = re.search(r"```json?\s*(.*?)\s*```", raw_text, re.DOTALL)
        if match:
            data = json.loads(match.group(1))
        else:
            raise ValueError(f"Groq returned invalid JSON: {e}\n\n{raw_text[:500]}")

    result = ScenarioGenerationResult.model_validate(data)

    # Enrich event_type and context safely
    for s in result.scenarios:
        # Enforce globally unique IDs to prevent database collisions
        s.id = f"SC-{uuid.uuid4().hex[:10].upper()}"
        
        if s.givens_table:
            if s.givens_table.event_type:
                s.event_type = s.givens_table.event_type
            else:
                s.givens_table.event_type = "normal"
                s.event_type = "normal"
                
            if not s.givens_table.context:
                s.givens_table.context = "General market activity"
        else:
            s.event_type = "normal"

    return result
