from typing import List, Any, Literal
from pydantic import BaseModel, ValidationError
from langchain_core.prompts import ChatPromptTemplate
import hashlib
import json
import re
from ..observability.logger import SessionLogger


# --- Entity Model ---
# Defines the structure for extracted financial entities
class Entity(BaseModel):
    name: str
    type: str  # company, index, sector, macro
    ticker: str | None = None


# --- NEW: Query Intent Model ---
# Classifies the type of financial question being asked
class QueryIntent(BaseModel):
    category: Literal["quantitative", "qualitative", "mixed"]
    reasoning: str = ""


# Combined extraction result
class ExtractionResult(BaseModel):
    entities: List[Entity]
    intent: QueryIntent


NER_PROMPT = """
Extract financial entities from the user query AND classify query intent.

Rules for entities:
- Identify companies, indices, sectors, macro indicators
- Include ticker if confidently known
- If uncertain, leave ticker null
- Do NOT guess

Rules for intent classification:
- "quantitative": Query asks for financial metrics, performance data, or forward-looking estimates
  Keywords: revenue, earnings, margins, growth rates, valuations, debt levels, price targets, 
  expected returns, potential gains, stock performance, earnings growth, profit margins, ROE, 
  cash flow, balance sheet metrics, or ANY metric/measurement request
- "qualitative": Query asks for explanations or narrative understanding
  Keywords: business model, strategy, competitive position, how they make money, operations,
  moat, management quality, industry dynamics
- "mixed": Query contains both types of questions

Examples:
- "What is Apple's revenue?" → quantitative
- "What is the growth potential for Nvidia?" → quantitative
- "How much increment in stocks can we expect?" → quantitative
- "What are Tesla's expected returns in 6 months?" → quantitative
- "Explain Uber's business model" → qualitative
- "How does Amazon make money?" → qualitative
- "Compare Tesla and Ford operating margins" → quantitative
- "What is Microsoft's strategy AND revenue growth?" → mixed

Return JSON only in this format:
{{
  "entities": [
    {{"name": "Company Name", "type": "company", "ticker": "TICKER"}}
  ],
  "intent": {{
    "category": "quantitative|qualitative|mixed",
    "reasoning": "Brief explanation of intent classification"
  }}
}}

Query:
{query}
"""


# --- Entity Extractor ---
# Handles the interpretation of user queries to identify financial entities and intent
class EntityExtractor:
    def __init__(self, llm: Any, logger: SessionLogger | None = None):
        # FIX 6: Verify LLM is deterministic
        if hasattr(llm, 'temperature') and llm.temperature != 0:
            raise ValueError(f"EntityExtractor requires deterministic LLM (temperature=0), got {llm.temperature}")
        
        self.llm = llm
        self.logger = logger or SessionLogger()

    def _extract_json(self, text: str) -> str:
        """Extract JSON from text, handling markdown code blocks."""
        # Remove markdown code block markers
        text = re.sub(r'```json\s*', '', text)
        text = re.sub(r'```\s*$', '', text)
        text = text.strip()
        
        # Find the first { and match it with the last }
        start_idx = text.find('{')
        if start_idx == -1:
            return text
        
        # Find matching closing brace by counting
        brace_count = 0
        for i in range(start_idx, len(text)):
            if text[i] == '{':
                brace_count += 1
            elif text[i] == '}':
                brace_count -= 1
                if brace_count == 0:
                    return text[start_idx:i+1]
        
        return text[start_idx:]

    async def extract(self, query: str) -> ExtractionResult:
        """
        Extract financial entities AND classify query intent from user query.
        
        Returns:
            ExtractionResult with both entities and intent classification
        """
        self.logger.log("ENTITY_EXTRACTION_STARTED", {"query": query})
        prompt = ChatPromptTemplate.from_template(NER_PROMPT)
        chain = prompt | self.llm

        data: dict | None = None
        last_err: Exception | None = None
        raw = ""
        for attempt in range(3):
            try:
                response = await chain.ainvoke({"query": query})
                raw = response.content
                json_text = self._extract_json(raw)
                data = json.loads(json_text)
                break
            except (json.JSONDecodeError, ValueError) as e:
                # Parse errors are deterministic at temperature=0 — do not retry.
                last_err = e
                self.logger.log(
                    "ENTITY_EXTRACTION_PARSE_FATAL",
                    {
                        "raw_hash": hashlib.sha256(raw.encode()).hexdigest(),
                        "raw_len": len(raw),
                        "raw_preview": raw[:200],
                        "error": str(e),
                    },
                )
                raise RuntimeError(
                    "Failed to parse entity extractor output as JSON"
                ) from e
            except Exception as e:
                # Transient network/API error — retry up to 3 times
                last_err = e
                self.logger.log(
                    "ENTITY_EXTRACTION_PARSE_ERROR",
                    {"attempt": attempt + 1, "error": str(e)},
                )

        if data is None:
            self.logger.log(
                "ENTITY_EXTRACTION_PARSE_FATAL",
                {"error": str(last_err)},
            )
            raise RuntimeError(
                "Entity extraction failed after 3 attempts (transient API error)"
            ) from last_err

        # Extract entities
        entities = []
        for e in data.get("entities", []):
            try:
                ent = Entity(**e)
                entities.append(ent)
            except ValidationError as ve:
                # skip invalid entities but log
                self.logger.log("ENTITY_VALIDATION_ERROR", {"entity": e, "error": ve.errors()})

        # Extract intent
        intent_data = data.get("intent", {})
        try:
            intent = QueryIntent(**intent_data)
        except (ValidationError, TypeError):
            # Fallback to quantitative if intent classification fails
            self.logger.log("INTENT_EXTRACTION_ERROR", {"intent_data": intent_data})
            intent = QueryIntent(category="quantitative", reasoning="Defaulted due to extraction error")

        result = ExtractionResult(entities=entities, intent=intent)
        self.logger.log("ENTITY_EXTRACTION_COMPLETED", {"count": len(entities), "intent": intent.category})
        return result
