from __future__ import annotations

from typing import Any, Dict, Iterable, Protocol

from analogai.deepthink.presentation.parsers.utils.relationship_type_normalizer import (
    tag_key_to_relationship_type,
)


def normalize_transformed_sentence(payload: Any) -> Dict[str, Any]:
    return TransformedSentenceNormalizer().normalize(payload)


class StructureRule(Protocol):
    def matches(self, payload: Any) -> bool:
        raise NotImplementedError

    def normalize(self, payload: Any) -> Dict[str, Any]:
        raise NotImplementedError


class _ListStructureRule:
    def matches(self, payload: Any) -> bool:
        return isinstance(payload, list)

    def normalize(self, payload: Any) -> Dict[str, Any]:
        return {"sentences": payload}


class _NonDictStructureRule:
    def matches(self, payload: Any) -> bool:
        return not isinstance(payload, dict)

    def normalize(self, payload: Any) -> Dict[str, Any]:
        return {"sentences": [payload]}


class _SentencesKeyRule:
    def matches(self, payload: Any) -> bool:
        return isinstance(payload, dict) and "sentences" in payload

    def normalize(self, payload: Any) -> Dict[str, Any]:
        return payload


class _RelationshipsKeyRule:
    def matches(self, payload: Any) -> bool:
        return isinstance(payload, dict) and "relationships" in payload

    def normalize(self, payload: Any) -> Dict[str, Any]:
        return {"sentences": payload["relationships"]}


class _IntentKeyRule:
    def matches(self, payload: Any) -> bool:
        return isinstance(payload, dict) and ("intent" in payload or "intent_type" in payload)

    def normalize(self, payload: Any) -> Dict[str, Any]:
        return {"sentences": [payload]}


class _DefaultStructureRule:
    def matches(self, payload: Any) -> bool:
        return True

    def normalize(self, payload: Any) -> Dict[str, Any]:
        return payload


class StructureNormalizer:
    def __init__(self, rules: Iterable[StructureRule] | None = None) -> None:
        self._rules = list(rules) if rules is not None else [
            _ListStructureRule(),
            _NonDictStructureRule(),
            _SentencesKeyRule(),
            _RelationshipsKeyRule(),
            _IntentKeyRule(),
            _DefaultStructureRule(),
        ]

    def normalize(self, payload: Any) -> Dict[str, Any]:
        for rule in self._rules:
            if rule.matches(payload):
                return rule.normalize(payload)
        return {"sentences": [payload]}


class SentenceTransformer(Protocol):
    def transform(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        raise NotImplementedError


class _PredicateTransformer:
    def transform(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        predicate = payload.get("predicate")
        if isinstance(predicate, dict):
            payload["predicate"] = self._normalize_predicate(predicate)
        return payload

    def _normalize_predicate(self, predicate: Dict[str, Any]) -> Dict[str, Any]:
        relation = predicate.get("relation")
        if relation:
            rel_type = tag_key_to_relationship_type(str(relation))
            if rel_type is not None:
                predicate["relation"] = rel_type.value
        return predicate


class _ClauseListTransformer:
    def __init__(self, keys: Iterable[str] | None = None) -> None:
        self._keys = list(keys) if keys is not None else ["causes", "effects"]

    def transform(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        for key in self._keys:
            clauses = payload.get(key)
            if isinstance(clauses, list):
                payload[key] = [self._normalize_clause(clause) for clause in clauses]
        return payload

    def _normalize_clause(self, clause: Any) -> Any:
        if isinstance(clause, dict):
            predicate = clause.get("predicate")
            if isinstance(predicate, dict):
                clause["predicate"] = _PredicateTransformer()._normalize_predicate(predicate)
        return clause


class RelationTypeNormalizer:
    def __init__(self, transformers: Iterable[SentenceTransformer] | None = None) -> None:
        self._transformers = list(transformers) if transformers is not None else [
            _PredicateTransformer(),
            _ClauseListTransformer(),
        ]

    def normalize_sentence(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        for transformer in self._transformers:
            payload = transformer.transform(payload)
        return payload


class TransformedSentenceNormalizer:
    def __init__(
        self,
        structure_normalizer: StructureNormalizer | None = None,
        relation_normalizer: RelationTypeNormalizer | None = None,
    ) -> None:
        self._structure = structure_normalizer or StructureNormalizer()
        self._relation = relation_normalizer or RelationTypeNormalizer()

    def normalize(self, payload: Any) -> Dict[str, Any]:
        normalized = self._structure.normalize(payload)
        sentences = normalized.get("sentences")
        if isinstance(sentences, list):
            normalized["sentences"] = [self._normalize_sentence(sentence) for sentence in sentences]
        return normalized

    def _normalize_sentence(self, sentence: Any) -> Any:
        if isinstance(sentence, dict):
            return self._relation.normalize_sentence(sentence)
        return sentence
