from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Sequence, Tuple, TypeVar, cast

from ragu.chunker.types import Chunk
from ragu.common.global_parameters import Settings
from ragu.common.logger import logger
from ragu.common.prompts.default_models import (
    DefinitionBasedEntityDefinitionModel,
    DefinitionBasedEntityExtractionModel,
    DefinitionBasedRelationDefinitionModel,
    DefinitionBasedRelationExtractionModel,
)
from ragu.common.prompts.messages import ChatMessages, render
from ragu.common.prompts.prompt_storage import RAGUInstruction
from ragu.graph.types import Entity, Relation
from ragu.models.llm import LLM
from ragu.triplet.base_artifact_extractor import BaseArtifactExtractor
from ragu.triplet.prompts import (
    DEFINITION_BASED_ENTITY_DEFINITION_INSTRUCTION,
    DEFINITION_BASED_ENTITY_EXTRACTION_INSTRUCTION,
    DEFINITION_BASED_RELATION_DEFINITION_INSTRUCTION,
    DEFINITION_BASED_RELATION_EXTRACTION_INSTRUCTION,
)

T = TypeVar(
    "T",
    DefinitionBasedEntityExtractionModel,
    DefinitionBasedEntityDefinitionModel,
    DefinitionBasedRelationExtractionModel,
    DefinitionBasedRelationDefinitionModel,
)


@dataclass
class DefinitionChunkContext:
    chunk: Chunk
    normalized_entities: List[str] = field(default_factory=list)
    entities: List[Entity] = field(default_factory=list)
    relation_candidate_pairs: List[Tuple[str, str]] = field(default_factory=list)
    relations: List[Relation] = field(default_factory=list)


class DefinitionBasedArtifactsExtractorLLM(BaseArtifactExtractor):
    """
    Extract entities and relations with prompt fan-out over entities and entity pairs.

    Pipeline:
      1. Extract normalized entities per chunk.
      2. Generate a contextual definition for each entity.
      3. Extract relation candidates for each entity against the full entity list.
      4. Generate a contextual definition for each unique candidate pair.
    """

    NEGATION_PATTERNS: Sequence[str] = (
        r"^\s*$",
        r"^\s*[\-–—]\s*$",
        r"^\s*отношения между этими сущностями нет\s*$",
        r"\bотношени\w*\s+между\s+этими\s+сущност\w*\s+нет\b",
        r"\bнет\s+(?:отношени\w*|связ\w*)\b",
        r"\bотсутств\w*\s+(?:отношени\w*|связ\w*)\b",
        r"\bне\s+содерж\w*[^.\n]{0,80}(?:отношени\w*|связ\w*)\b",
    )

    def __init__(
        self,
        llm: LLM,
        language: str | None = None,
        temperature: float = 0.0,
        top_p: float = 0.95,
        default_entity_type: str = "UNKNOWN",
        default_relation_type: str = "RELATED_TO",
        default_relation_strength: float = 1.0,
    ) -> None:
        prompts = {
            "entity_extraction": DEFINITION_BASED_ENTITY_EXTRACTION_INSTRUCTION,
            "entity_definition": DEFINITION_BASED_ENTITY_DEFINITION_INSTRUCTION,
            "relation_extraction": DEFINITION_BASED_RELATION_EXTRACTION_INSTRUCTION,
            "relation_definition": DEFINITION_BASED_RELATION_DEFINITION_INSTRUCTION,
        }
        super().__init__(prompts=prompts)

        self.llm = llm
        self.language = language if language else Settings.language
        self.temperature = temperature
        self.top_p = top_p
        self.default_entity_type = default_entity_type
        self.default_relation_type = default_relation_type
        self.default_relation_strength = default_relation_strength
        self._negative_relation_re = re.compile(
            r"(?:" + "|".join(self.NEGATION_PATTERNS) + r")",
            flags=re.IGNORECASE | re.UNICODE,
        )

    async def extract(
        self,
        chunks: List[Chunk],
        *args: Any,
        **kwargs: Any,
    ) -> Tuple[List[Entity], List[Relation]]:
        if not chunks:
            return [], []

        contexts = [DefinitionChunkContext(chunk=chunk) for chunk in chunks]

        logger.info(f"Stage 1/4: Extracting normalized entities from {len(chunks)} chunks...")
        await self._extract_entities(contexts)

        total_entities = sum(len(ctx.normalized_entities) for ctx in contexts)
        logger.info(f"Stage 2/4: Generating entity definitions for {total_entities} entities...")
        await self._define_entities(contexts)

        logger.info("Stage 3/4: Extracting relation candidates...")
        await self._extract_relation_candidates(contexts)

        total_pairs = sum(len(ctx.relation_candidate_pairs) for ctx in contexts)
        logger.info(f"Stage 4/4: Generating relation definitions for {total_pairs} pairs...")
        await self._define_relations(contexts)

        entities = [entity for ctx in contexts for entity in ctx.entities]
        relations = [relation for ctx in contexts for relation in ctx.relations]
        return entities, relations

    async def _extract_entities(self, contexts: List[DefinitionChunkContext]) -> None:
        instruction: RAGUInstruction = self.get_prompt("entity_extraction")
        assert instruction.pydantic_model is DefinitionBasedEntityExtractionModel
        conversations: List[ChatMessages] = render(
            instruction.messages,
            source_text=[ctx.chunk.content for ctx in contexts],
        )
        responses = await self._run(
            conversations,
            output_schema=DefinitionBasedEntityExtractionModel,
            description="Definition-based entity extraction",
        )

        for ctx, response in zip(contexts, responses):
            ctx.normalized_entities = self._deduplicate_names(
                [entity.entity_name for entity in response.entities]
            )
            logger.debug(
                f"Chunk {ctx.chunk.id}: extracted {len(ctx.normalized_entities)} normalized entities"
            )

    async def _define_entities(self, contexts: List[DefinitionChunkContext]) -> None:
        instruction: RAGUInstruction = self.get_prompt("entity_definition")
        assert instruction.pydantic_model is DefinitionBasedEntityDefinitionModel

        conversations: List[ChatMessages] = []
        prompt_map: List[Tuple[DefinitionChunkContext, str]] = []

        for ctx in contexts:
            if not ctx.normalized_entities:
                continue

            chunk_conversations: List[ChatMessages] = render(
                instruction.messages,
                normalized_entity_in_text=ctx.normalized_entities,
                source_text=ctx.chunk.content,
            )
            for conversation, entity_name in zip(chunk_conversations, ctx.normalized_entities):
                conversations.append(conversation)
                prompt_map.append((ctx, entity_name))

        if not conversations:
            return

        responses = await self._run(
            conversations,
            output_schema=DefinitionBasedEntityDefinitionModel,
            description="Definition-based entity definitions",
        )

        for (ctx, entity_name), response in zip(prompt_map, responses):
            normalized_name = self._normalize_name_for_lookup(response.entity_name) or entity_name
            description = self._clean_text_response(response.description)
            if not description:
                continue

            ctx.entities.append(
                Entity(
                    entity_name=normalized_name,
                    entity_type=self.default_entity_type,
                    description=description,
                    source_chunk_id=[ctx.chunk.id],
                    documents_id=[ctx.chunk.doc_id] if ctx.chunk.doc_id else [],
                    clusters=[],
                )
            )

    async def _extract_relation_candidates(self, contexts: List[DefinitionChunkContext]) -> None:
        instruction: RAGUInstruction = self.get_prompt("relation_extraction")
        assert instruction.pydantic_model is DefinitionBasedRelationExtractionModel

        conversations: List[ChatMessages] = []
        prompt_map: List[Tuple[DefinitionChunkContext, str]] = []

        for ctx in contexts:
            if len(ctx.entities) < 2:
                continue

            entity_names = [entity.entity_name for entity in ctx.entities]
            chunk_conversations: List[ChatMessages] = render(
                instruction.messages,
                source_text=ctx.chunk.content,
                full_list_of_normalized_entities="\n".join(entity_names),
                first_entity_in_text=entity_names,
            )
            for conversation, entity_name in zip(chunk_conversations, entity_names):
                conversations.append(conversation)
                prompt_map.append((ctx, entity_name))

        if not conversations:
            return

        responses = await self._run(
            conversations,
            output_schema=DefinitionBasedRelationExtractionModel,
            description="Definition-based relation candidates",
        )

        candidate_pairs_by_context: Dict[int, List[Tuple[str, str]]] = {
            id(ctx): [] for ctx in contexts
        }
        seen_pairs_by_context: Dict[int, set[Tuple[str, str]]] = {
            id(ctx): set() for ctx in contexts
        }

        for (ctx, source_name), response in zip(prompt_map, responses):
            resolved_names = self._resolve_relation_targets(
                ctx.entities,
                [entity.entity_name for entity in response.related_entities],
            )
            for target_name in resolved_names:
                pair = self._canonical_pair(source_name, target_name)
                if pair[0] == pair[1]:
                    continue
                if pair in seen_pairs_by_context[id(ctx)]:
                    continue
                seen_pairs_by_context[id(ctx)].add(pair)
                candidate_pairs_by_context[id(ctx)].append(pair)

        for ctx in contexts:
            ctx.relation_candidate_pairs = candidate_pairs_by_context[id(ctx)]
            logger.debug(
                f"Chunk {ctx.chunk.id}: extracted {len(ctx.relation_candidate_pairs)} unique relation candidates"
            )

    async def _define_relations(self, contexts: List[DefinitionChunkContext]) -> None:
        instruction: RAGUInstruction = self.get_prompt("relation_definition")
        assert instruction.pydantic_model is DefinitionBasedRelationDefinitionModel

        conversations: List[ChatMessages] = []
        prompt_map: List[Tuple[DefinitionChunkContext, str, str]] = []

        for ctx in contexts:
            if not ctx.relation_candidate_pairs:
                continue

            first_entities = [pair[0] for pair in ctx.relation_candidate_pairs]
            second_entities = [pair[1] for pair in ctx.relation_candidate_pairs]
            chunk_conversations: List[ChatMessages] = render(
                instruction.messages,
                first_entity_in_text=first_entities,
                second_entity_in_text=second_entities,
                source_text=ctx.chunk.content,
            )
            for conversation, pair in zip(chunk_conversations, ctx.relation_candidate_pairs):
                conversations.append(conversation)
                prompt_map.append((ctx, pair[0], pair[1]))

        if not conversations:
            return

        responses = await self._run(
            conversations,
            output_schema=DefinitionBasedRelationDefinitionModel,
            description="Definition-based relation definitions",
        )

        for (ctx, subject_name, object_name), response in zip(prompt_map, responses):
            subject_name = self._normalize_name_for_lookup(response.first_entity) or subject_name
            object_name = self._normalize_name_for_lookup(response.second_entity) or object_name
            description = self._clean_relation_description(response.description)
            if not response.has_relation or not description or self._is_negative_relation(description):
                continue

            entity_by_name = {entity.entity_name: entity for entity in ctx.entities}
            subject_entity = entity_by_name.get(subject_name)
            object_entity = entity_by_name.get(object_name)
            if not subject_entity or not object_entity:
                logger.debug(
                    "Skipping relation with unresolved endpoints: "
                    f"{subject_name} -> {object_name}"
                )
                continue

            ctx.relations.append(
                Relation(
                    subject_id=subject_entity.id,
                    object_id=object_entity.id,
                    subject_name=subject_entity.entity_name,
                    object_name=object_entity.entity_name,
                    relation_type=self.default_relation_type,
                    description=description,
                    relation_strength=self.default_relation_strength,
                    source_chunk_id=[ctx.chunk.id],
                )
            )

    async def _run(
        self,
        conversations: List[ChatMessages],
        output_schema: type[T],
        description: str,
    ) -> List[T]:
        if not conversations:
            return []

        responses = await self.llm.batch_chat_completion(
            [conversation.to_openai() for conversation in conversations],
            output_schema=output_schema,
            temperature=self.temperature,
            top_p=self.top_p,
            desc=description,
        )
        return cast(List[T], list(responses))

    @staticmethod
    def _canonical_pair(left: str, right: str) -> Tuple[str, str]:
        if left.casefold() <= right.casefold():
            return left, right
        return right, left

    @staticmethod
    def _clean_text_response(response: str) -> str:
        cleaned = response.strip().strip("`").strip()
        return re.sub(r"\s+", " ", cleaned)

    @classmethod
    def _clean_relation_description(cls, response: str) -> str:
        cleaned = cls._clean_text_response(response)
        cleaned = re.sub(r"^[\-\u2022]\s*", "", cleaned)
        return cleaned

    @staticmethod
    def _deduplicate_names(items: List[str]) -> List[str]:
        return list(dict.fromkeys(items))

    @classmethod
    def _normalize_name_for_lookup(cls, value: str) -> str:
        cleaned = cls._clean_text_response(value)
        cleaned = cleaned.strip("\"'")
        cleaned = cleaned.rstrip(".,;:")
        return cleaned

    @classmethod
    def _resolve_relation_targets(
        cls,
        entities: List[Entity],
        candidates: List[str],
    ) -> List[str]:
        direct_lookup = {entity.entity_name: entity.entity_name for entity in entities}
        folded_lookup = {
            entity.entity_name.casefold(): entity.entity_name for entity in entities
        }

        resolved: List[str] = []
        for candidate in candidates:
            normalized = cls._normalize_name_for_lookup(candidate)
            if normalized in direct_lookup:
                resolved.append(direct_lookup[normalized])
                continue

            folded = folded_lookup.get(normalized.casefold())
            if folded:
                resolved.append(folded)

        return list(dict.fromkeys(resolved))

    def _is_negative_relation(self, description: str) -> bool:
        return bool(self._negative_relation_re.search(description))
