from __future__ import annotations

import json
import re
from typing import Iterable, List, Optional

from analogai.foundation.common.logging.app_logger import app_logger
from analogai.deepthink.infrastructure.gateways.llm.llm_factory import LlmGatewayFactory
from analogai.deepthink.infrastructure.gateways.llm.llm_gateway_provider_factory import LlmGatewayProviderFactory


_DEFAULT_STOPWORDS = {
    "a", "an", "and", "are", "as", "at", "be", "by", "can", "could", "did", "do", "does",
    "for", "from", "had", "has", "have", "how", "i", "if", "in", "is", "it", "may", "might",
    "of", "on", "or", "should", "so", "that", "the", "their", "them", "then", "there", "these",
    "they", "this", "to", "was", "were", "what", "when", "where", "which", "who", "why", "will",
    "with", "would", "you", "your",
    "sentences", "intent_type", "intent", "subject", "predicate", "relation", "object", "complement",
    "wrestler", "show", "event", "events", "match", "matches", "tournament",
    "film", "films", "movie", "movies", "show", "series", "episode", "episodes",
}


class QuestionEntityExtractor:
    def __init__(self, llm_gateway=None, max_llm_tokens: int = 600):
        self._llm_gateway = llm_gateway
        self._max_llm_tokens = max_llm_tokens

    def extract(self, text: str, max_entities: int = 8) -> List[str]:
        if not text:
            return []
        candidates = []
        candidates.extend(self._extract_with_llm(text))
        candidates.extend(_fallback_tokens(text))

        normalized = _dedupe_preserve_order(_normalize_entities(candidates))
        return normalized[:max_entities]

    def _extract_with_llm(self, text: str) -> List[str]:
        gateway = self._llm_gateway or LlmGatewayProviderFactory.create_gateway(
            system_prompt="",
            max_tokens=self._max_llm_tokens,
        )
        sanitized = text.replace('"', "")
        prompts = [
            (
                "You extract named entities from a question. Output ONLY valid JSON (no markdown, no prose). "
                "Schema: {\"entities\": [\"...\", \"...\"]}. Use double quotes; no trailing commas. "
                "\n\nTHINKING (do not output):\n"
                "1) Identify only proper names and titled works.\n"
                "2) Exclude numbers, years, dates, ages, quantities.\n"
                "3) Exclude roles/titles/occupations and generic nouns (e.g., wrestler, director, show, album, movie, event, city, country).\n"
                "4) Keep only the named entity strings.\n\n"
                "DO NOT return: roles, descriptors, generic nouns, adjectives, or numeric tokens.\n"
                "DO NOT include: subjects/relations/predicates/certainty or any extra fields.\n\n"
                "Examples:\n"
                "Q: Which wrestler, who died in 2009, featured on the wrestling show Starrcade?\n"
                "A: {\"entities\": [\"Starrcade\"]}\n"
                "Q: The song Medicine by The 1975 appears in a film directed by what Danish filmmaker?\n"
                "A: {\"entities\": [\"Medicine\", \"The 1975\"]}\n"
                "Q: World for Ransom was directed by the producer notable for which 1974 film?\n"
                "A: {\"entities\": [\"World for Ransom\"]}\n\n"
                f"Sentence: {sanitized}"
            ),
            (
                "Return ONLY valid JSON with double quotes and no trailing commas. "
                "Schema: {\"entities\": [\"...\", \"...\"]}. "
                "Only named entities (proper nouns / titled works). "
                "Exclude years, dates, numbers, roles/titles/occupations, and generic nouns. "
                "Examples: 'Danish filmmaker' -> exclude; '1974' -> exclude; 'director' -> exclude. "
                f"Sentence: {sanitized}"
            ),
        ]

        last_response = ""
        for attempt, user_message in enumerate(prompts, start=1):
            try:
                last_response = gateway.generate_completion(user_message)
            except Exception as exc:
                app_logger.warning(
                    "[QuestionEntityExtractor] LLM extraction failed. "
                    f"attempt={attempt} reason={exc} text={text}"
                )
                continue

            try:
                data = _safe_json_loads(last_response)
                return _extract_entities_from_json(data)
            except json.JSONDecodeError as exc:
                app_logger.warning(
                    "[QuestionEntityExtractor] Invalid JSON from LLM. "
                    f"attempt={attempt} reason={exc} text={text} llm_output={last_response}"
                )

        return []


def _fallback_tokens(text: str) -> Iterable[str]:
    words = re.findall(r"[A-Za-z][A-Za-z0-9'_-]*", text)
    tokens: List[str] = []
    for word in words:
        lower = word.lower()
        if lower in _DEFAULT_STOPWORDS:
            continue
        if len(lower) <= 2:
            continue
        tokens.append(word)
    tokens.extend(_extract_years(text))
    return tokens


def _extract_years(text: str) -> List[str]:
    years = re.findall(r"\b(18\d{2}|19\d{2}|20\d{2})\b", text)
    return list(dict.fromkeys(years))


def _normalize_entities(entities: Iterable[str]) -> List[str]:
    normalized: List[str] = []
    for entity in entities:
        cleaned = re.sub(r"\s+", " ", entity).strip(" \t\n\r\f\v.,!?;:\"'()[]{}")
        if not cleaned:
            continue
        normalized.append(cleaned)
    return normalized


def _dedupe_preserve_order(items: Iterable[str]) -> List[str]:
    seen = set()
    result = []
    for item in items:
        key = item.lower()
        if key in seen:
            continue
        seen.add(key)
        result.append(item)
    return result


def _extract_entities_from_json(data) -> List[str]:
    if not isinstance(data, dict):
        return []
    entities_field = data.get("entities")
    if isinstance(entities_field, list):
        entities: List[str] = []
        for item in entities_field:
            if isinstance(item, str):
                entities.append(item)
        return entities
    sentences = data.get("sentences")
    if isinstance(sentences, list):
        items = []
        for sentence in sentences:
            items.extend(_extract_entities_from_json(sentence))
        return items

    entities: List[str] = []
    entities.extend(_extract_entities_from_node(data.get("subject")))
    predicate = data.get("predicate")
    if isinstance(predicate, dict):
        entities.extend(_extract_entities_from_node(predicate.get("complement")))

    return entities


def _safe_json_loads(text: str):
    cleaned = text.strip()
    if not cleaned:
        raise json.JSONDecodeError("Empty JSON", cleaned, 0)

    if cleaned.startswith("```"):
        cleaned = re.sub(r"^```[a-zA-Z]*\n?", "", cleaned)
        cleaned = re.sub(r"```$", "", cleaned).strip()

    try:
        return json.loads(cleaned)
    except json.JSONDecodeError:
        pass

    start = cleaned.find("{")
    end = cleaned.rfind("}")
    if start != -1 and end != -1 and end > start:
        snippet = cleaned[start:end + 1]
        return json.loads(snippet)

    return json.loads(cleaned)


def _extract_entities_from_node(node: Optional[object]) -> List[str]:
    if node is None:
        return []
    if isinstance(node, str):
        return [node]
    if isinstance(node, dict):
        caption = node.get("caption")
        entities = [caption] if isinstance(caption, str) else []
        attributes = node.get("attributes")
        if isinstance(attributes, list):
            for attr in attributes:
                if not isinstance(attr, dict):
                    continue
                value = attr.get("value")
                if isinstance(value, str):
                    entities.append(value)
        return entities
    return []
