from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Iterable, List, Optional, Protocol

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


_RELATION_TERMS: list[str] = []


class EntityExtractionService:
    def __init__(self, llm_gateway=None, max_llm_tokens: int = 600):
        self._llm_gateway = llm_gateway
        self._max_llm_tokens = max_llm_tokens
        self._schema_hint = _load_prompt_schema("question_entities.schema.json")

    def extract_entities_and_relations(
        self,
        text: str,
        max_entities: int = 8,
        max_relations: int = 6,
    ) -> tuple[List[str], List[str]]:
        if not text:
            return [], []
        llm_entities = self._extract_with_llm(text)
        normalized_entities = _dedupe_preserve_order(_normalize_entities(llm_entities))
        limited_entities = normalized_entities[:max_entities]
        limited_relations: List[str] = []
        app_logger.info(
            "[EntityExtractionService] entities=%s",
            limited_entities,
        )
        return limited_entities, limited_relations

    def extract_keywords(self, text: str, max_keywords: int = 8) -> List[str]:
        entities, _ = self.extract_entities_and_relations(text)
        combined = _dedupe_preserve_order(_normalize_entities(list(entities)))
        return combined[:max_keywords]

    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 = text.replace('"', "")
        prompt_variants = [
            (
                "You extract keywords and relation verbs from a question. Output ONLY valid JSON (no markdown, no prose). "
                f"Schema: {self._schema_hint}. Use double quotes; no trailing commas. "
                "\n\nTHINKING (do not output):\n"
                "1) Identify proper names and titled works.\n"
                "2) Include technical identifiers and acronyms (e.g., llm_model, LLM, API keys, config names).\n"
                "3) Extract noun phrases, relation verbs, and key nouns relevant to the question (e.g., 'llm model', 'gateway implementation', 'contact', 'founded by').\n"
                "4) Exclude numbers, years, dates, ages, quantities unless they are part of a technical identifier.\n"
                "5) Exclude roles/titles/occupations and generic filler nouns unless they are the subject of the question.\n"
                "6) Keep only named entities, relation verbs, noun phrases, key nouns, or technical identifiers.\n\n"
                "DO NOT return: adjectives alone or numeric tokens (unless part of a technical identifier).\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: What is llm_model?\n"
                "A: {\"entities\": [\"llm_model\"]}\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_text}"
            ),
            (
                "Return ONLY valid JSON with double quotes and no trailing commas. "
                f"Schema: {self._schema_hint}. "
                "Include named entities (proper nouns / titled works), technical identifiers (snake_case, acronyms, config names), "
                "relation verbs, and relevant noun phrases or key nouns from the question. "
                "Exclude years, dates, numbers, roles/titles/occupations, and generic filler nouns unless part of a technical identifier. "
                "Examples: 'Danish filmmaker' -> exclude; '1974' -> exclude; 'director' -> exclude; 'llm_model' -> include; 'llm model' -> include. "
                f"Sentence: {sanitized_text}"
            ),
        ]

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

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

        return []


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


class JsonEntityExtractor(Protocol):
    def can_handle(self, payload: dict) -> bool:
        ...

    def extract(self, payload: dict) -> List[str]:
        ...


class EntitiesListExtractor:
    def can_handle(self, payload: dict) -> bool:
        return isinstance(payload.get("entities"), list)

    def extract(self, payload: dict) -> List[str]:
        entities_field = payload.get("entities")
        if not isinstance(entities_field, list):
            return []
        return [item for item in entities_field if isinstance(item, str)]


class SentencesExtractor:
    def can_handle(self, payload: dict) -> bool:
        return isinstance(payload.get("sentences"), list)

    def extract(self, payload: dict) -> List[str]:
        sentences = payload.get("sentences")
        if not isinstance(sentences, list):
            return []
        extracted: List[str] = []
        for sentence in sentences:
            extracted.extend(_extract_entities_from_json(sentence))
        return extracted


class SubjectPredicateExtractor:
    def can_handle(self, payload: dict) -> bool:
        return "subject" in payload or "predicate" in payload

    def extract(self, payload: dict) -> List[str]:
        entities: List[str] = []
        entities.extend(_extract_entities_from_node(payload.get("subject")))
        predicate = payload.get("predicate")
        if isinstance(predicate, dict):
            entities.extend(_extract_entities_from_node(predicate.get("complement")))
        return entities


def _json_entity_extractors() -> list[JsonEntityExtractor]:
    return [
        EntitiesListExtractor(),
        SentencesExtractor(),
        SubjectPredicateExtractor(),
    ]


def _extract_entities_from_json(payload) -> List[str]:
    if not isinstance(payload, dict):
        return []
    for extractor in _json_entity_extractors():
        if not extractor.can_handle(payload):
            continue
        entities = extractor.extract(payload)
        if entities:
            return entities
    return []


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:
        repaired = _repair_json_fragment(cleaned)
        if repaired and repaired != cleaned:
            return json.loads(repaired)

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

    repaired = _repair_json_fragment(cleaned)
    if repaired and repaired != cleaned:
        return json.loads(repaired)
    return json.loads(cleaned)


def _repair_json_fragment(text: str) -> str:
    if not text:
        return text
    cleaned = text.strip()
    if not cleaned.startswith("{"):
        return cleaned
    if cleaned.endswith("}"):
        return cleaned
    brace_open = cleaned.count("{")
    brace_close = cleaned.count("}")
    bracket_open = cleaned.count("[")
    bracket_close = cleaned.count("]")
    missing_brackets = max(0, bracket_open - bracket_close)
    missing_braces = max(0, brace_open - brace_close)
    if missing_brackets == 0 and missing_braces == 0:
        return cleaned
    return f"{cleaned}{']' * missing_brackets}{'}' * missing_braces}"


def _trim_for_log(text: str, max_chars: int = 400) -> str:
    if text is None:
        return ""
    cleaned = " ".join(str(text).split())
    if len(cleaned) <= max_chars:
        return cleaned
    return f"{cleaned[: max_chars - 1]}…"


def _load_prompt_schema(schema_name: str) -> str:
    prompts_dir = Path(__file__).resolve().parents[2] / "infrastructure" / "prompts" / "schemas"
    schema_path = prompts_dir / schema_name
    if not schema_path.exists():
        raise FileNotFoundError(f"Schema file not found: {schema_path}")
    schema_text = schema_path.read_text(encoding="utf-8")
    try:
        schema_json = json.loads(schema_text)
        return json.dumps(schema_json, separators=(",", ":"))
    except json.JSONDecodeError:
        return schema_text.strip()


class NodeEntityExtractor(Protocol):
    def can_handle(self, node: Optional[object]) -> bool:
        ...

    def extract(self, node: Optional[object]) -> List[str]:
        ...


class StringNodeExtractor:
    def can_handle(self, node: Optional[object]) -> bool:
        return isinstance(node, str)

    def extract(self, node: Optional[object]) -> List[str]:
        return [node] if isinstance(node, str) else []


class DictNodeExtractor:
    def can_handle(self, node: Optional[object]) -> bool:
        return isinstance(node, dict)

    def extract(self, node: Optional[object]) -> List[str]:
        if not isinstance(node, dict):
            return []
        entities: List[str] = []
        caption = node.get("caption")
        if isinstance(caption, str):
            entities.append(caption)
        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


def _node_entity_extractors() -> list[NodeEntityExtractor]:
    return [
        StringNodeExtractor(),
        DictNodeExtractor(),
    ]


def _extract_entities_from_node(node: Optional[object]) -> List[str]:
    for extractor in _node_entity_extractors():
        if not extractor.can_handle(node):
            continue
        return extractor.extract(node)
    return []


def _extract_relations_from_text(text: str) -> List[str]:
    return []


# Backward compatibility
class QuestionEntityExtractionService(EntityExtractionService):
    def extract_entities(self, text: str, max_entities: int = 8) -> List[str]:
        entities, _ = super().extract_entities_and_relations(
            text,
            max_entities=max_entities,
        )
        return entities


    # Backward compatibility
    KeywordExtractionService = EntityExtractionService
