import logging
from collections import defaultdict
from collections.abc import Iterator
from dataclasses import dataclass, field
from typing import Any

import requests

from gocam.datamodel import (
    Activity,
    BiologicalProcessAssociation,
    CausalAssociation,
    CellTypeAssociation,
    CellularAnatomicalEntityAssociation,
    EnabledByAssociation,
    EnabledByGeneProductAssociation,
    EnabledByProteinComplexAssociation,
    EvidenceItem,
    Model,
    ModelStateEnum,
    MolecularFunctionAssociation,
    MoleculeAssociation,
    MoleculeNode,
    Object,
    PhaseAssociation,
    ProteinComplexMemberAssociation,
    ProvenanceInfo,
)
from gocam.translation.result import TranslationResult, TranslationWarning, WarningType
from gocam.vocabulary.taxa import TaxonVocabulary

ENABLED_BY = "RO:0002333"
PART_OF = "BFO:0000050"
HAS_PART = "BFO:0000051"
OCCURS_IN = "BFO:0000066"
HAS_INPUT = "RO:0002233"
HAS_OUTPUT = "RO:0002234"
HAS_PRIMARY_INPUT = "RO:0004009"
HAS_PRIMARY_OUTPUT = "RO:0004008"
HAPPENS_DURING = "RO:0002092"
LOCATED_IN = "RO:0001025"

logger = logging.getLogger(__name__)


def _normalize_property(prop: str) -> str:
    """
    Normalize a property.

    Sometimes the JSON will use full URIs, sometimes just the local part

    Args:
        prop: The property to normalize

    Returns:
        The normalized property
    """
    if "/" in prop:
        return prop.split("/")[-1]
    return prop


def _annotations(obj: dict) -> dict[str, str]:
    """
    Extract annotations from an object (assumes single-valued).

    Annotations are lists of objects with keys "key" and "value".

    Args:
        obj: The object to extract annotations from

    Returns:
        The extracted annotations
    """
    return {
        _normalize_property(a["key"]): a["value"] for a in obj.get("annotations", [])
    }


def _annotations_multivalued(obj: dict) -> dict[str, list[str]]:
    """
    Extract annotations from an object (assumes multi-valued).

    Annotations are lists of objects with keys "key" and "value".

    Args:
        obj: The object to extract annotations from

    Returns:
        The extracted annotations
    """
    anns = defaultdict(list)
    for a in obj.get("annotations", []):
        key = _normalize_property(a["key"])
        value = a["value"]
        anns[key].append(value)
    return anns


MOLECULAR_FUNCTION = "GO:0003674"
BIOLOGICAL_PROCESS = "GO:0008150"
CELLULAR_COMPONENT = "GO:0005575"
INFORMATION_BIOMACROMOLECULE = "CHEBI:33695"
PROTEIN_CONTAINING_COMPLEX = "GO:0032991"
EVIDENCE = "ECO:0000000"
CHEMICAL_ENTITY = "CHEBI:24431"
ANATOMICAL_ENTITY = "UBERON:0001062"


@dataclass
class MinervaWrapper:
    """
    An Wrapper over the current GO API which returns "Minerva" JSON objects.

    TODO: Implement a fact counter to ensure all facts are encountered for and nothing dropped on floor
    """

    session: requests.Session = field(default_factory=lambda: requests.Session())
    gocam_index_url: str = "https://go-public.s3.amazonaws.com/files/gocam-models.json"
    gocam_endpoint_base: str = "https://api.geneontology.org/api/go-cam/"

    def models(self) -> Iterator[Model]:
        """Iterator over all GO-CAM models from the index.

        This method fetches the list of all GO-CAM models from the index URL. For each model, the
        Minerva JSON object is fetched and converted to a Model object.

        Yields:
            GO-CAM Model
        """

        for gocam_id in self.models_ids():
            yield self.fetch_model(gocam_id)

    def models_ids(self) -> Iterator[str]:
        """Iterator over all GO-CAM IDs from the index.

        This method fetches the list of all GO-CAM models from the index URL and returns an
        iterator over the IDs of each model.

        Yields:
            GO-CAM ID
        """

        response = self.session.get(self.gocam_index_url)
        response.raise_for_status()
        for model in response.json():
            gocam = model.get("gocam")
            if gocam is None:
                raise ValueError(f"Missing gocam in {model}")
            yield gocam.replace("http://model.geneontology.org/", "")

    def fetch_minerva_object(self, gocam_id: str) -> dict:
        """Fetch a Minerva JSON object for a given GO-CAM ID.

        Args:
            gocam_id: GO-CAM ID

        Returns:
            Minerva JSON object
        """
        if not gocam_id:
            raise ValueError(f"Missing GO-CAM ID: {gocam_id}")
        local_id = gocam_id.replace("gocam:", "")
        url = f"{self.gocam_endpoint_base}{local_id}"
        response = self.session.get(url)
        response.raise_for_status()
        return response.json()

    def fetch_model(self, gocam_id: str) -> Model:
        """Fetch a GO-CAM Model for a given GO-CAM ID.

        Args:
            gocam_id: GO-CAM ID

        Returns:
            GO-CAM Model
        """
        minerva_object = self.fetch_minerva_object(gocam_id)
        return self.minerva_object_to_model(minerva_object)

    @staticmethod
    def translate(minerva_obj: dict) -> TranslationResult[Model]:
        """Convert a Minerva JSON object to a GO-CAM Model.

        Args:
            obj: Minerva JSON object

        Returns:
            Object containing GO-CAM Model and any translation warnings
        """
        id = minerva_obj["id"]

        # Bookkeeping variables

        # individual ID to "root" type / category, e.g Evidence, BP
        individual_to_root_types: dict[str, list[str]] = {}
        individual_to_term: dict[str, str] = {}
        individual_to_annotations: dict[str, dict[str, str]] = {}
        individual_to_annotations_multivalued: dict[str, dict[str, list[str]]] = {}
        objects_by_id: dict[str, dict] = {}
        activities: list[Activity] = []
        activities_by_mf_id: defaultdict[str, list[Activity]] = defaultdict(list)
        molecule_nodes_by_id: dict[str, MoleculeNode] = {}
        facts_by_property: defaultdict[str, list[dict]] = defaultdict(list)
        facts_by_subject_property: defaultdict[tuple[str, str], list[dict]] = (
            defaultdict(list)
        )
        processed_facts: set[tuple[str, str, str]] = set()

        translation_warnings: set[TranslationWarning] = set()

        def _setattr_with_warning(obj: Activity, attr: str, value: Any):
            """Set an attribute on an object, with a warning if it already exists.

            Args:
                obj: The Activity to set the attribute on
                attr: The attribute to set
                value: The value to set
            """
            if getattr(obj, attr, None) is not None:
                translation_warnings.add(
                    TranslationWarning(
                        type=WarningType.ATTRIBUTE_OVERWRITE,
                        message=f"Overwriting {attr} for {obj.id}",
                        entity_id=obj.id,
                    )
                )
            setattr(obj, attr, value)

        def _process_fact(fact: dict) -> tuple[list[EvidenceItem], ProvenanceInfo]:
            """Process a fact to extract evidence and provenance information."""
            processed_facts.add((fact["subject"], fact["property"], fact["object"]))

            annotations = _annotations(fact)
            annotations_mv = _annotations_multivalued(fact)

            evidence_inst_ids = annotations_mv.get("evidence", [])
            evs: list[EvidenceItem] = []
            for evidence_inst_id in evidence_inst_ids:
                evidence_inst_annotations = individual_to_annotations.get(
                    evidence_inst_id, {}
                )
                evidence_inst_annotations_multivalued = (
                    individual_to_annotations_multivalued.get(evidence_inst_id, {})
                )
                with_obj: str | None = evidence_inst_annotations.get("with", None)
                if with_obj:
                    with_objs = [s.strip() for s in with_obj.split("|")]
                else:
                    with_objs = None
                prov = ProvenanceInfo(
                    contributor=evidence_inst_annotations_multivalued.get(
                        "contributor"
                    ),
                    date=evidence_inst_annotations.get("date", None),
                    provided_by=evidence_inst_annotations_multivalued.get("providedBy"),
                )
                ev = EvidenceItem(
                    term=individual_to_term.get(evidence_inst_id, None),
                    reference=evidence_inst_annotations.get("source", None),
                    with_objects=with_objs,
                    provenances=[prov],
                )
                evs.append(ev)

            prov = ProvenanceInfo(
                contributor=annotations_mv.get("contributor"),
                date=annotations.get("date", None),
                provided_by=annotations_mv.get("providedBy"),
            )
            return evs, prov

        def _iter_activities_by_fact_subject(
            *,
            fact_property: str,
        ) -> Iterator[tuple[Activity, str, list[EvidenceItem], ProvenanceInfo]]:
            for fact in facts_by_property.get(fact_property, []):
                subject, object_ = fact["subject"], fact["object"]
                if object_ not in individual_to_term:
                    translation_warnings.add(
                        TranslationWarning(
                            type=WarningType.MISSING_TERM,
                            message=f"Missing term for object {object_} in fact",
                            entity_id=object_,
                        )
                    )
                    continue
                for activity in activities_by_mf_id.get(subject, []):
                    evs, provenance = _process_fact(fact)
                    yield activity, object_, evs, provenance

        def _add_molecule_node(individual_id: str) -> None:
            if individual_id in molecule_nodes_by_id:
                # Molecule node already exists, nothing to do
                return

            if individual_id not in individual_to_term:
                translation_warnings.add(
                    TranslationWarning(
                        type=WarningType.MISSING_TERM,
                        message=f"Missing term for individual {individual_id} when creating molecule node",
                        entity_id=individual_id,
                    )
                )
                return

            molecule_node = MoleculeNode(
                id=individual_id,
                term=individual_to_term[individual_id],
            )

            for located_in_fact in facts_by_subject_property.get(
                (individual_id, LOCATED_IN), []
            ):
                evs, prov = _process_fact(located_in_fact)
                if molecule_node.located_in is not None:
                    translation_warnings.add(
                        TranslationWarning(
                            type=WarningType.ATTRIBUTE_OVERWRITE,
                            message=f"Overwriting located_in for molecule {individual_id}",
                            entity_id=individual_id,
                        )
                    )
                molecule_node.located_in = CellularAnatomicalEntityAssociation(
                    term=individual_to_term.get(located_in_fact["object"], None),
                    evidence=evs,
                    provenances=[prov],
                )

            molecule_nodes_by_id[individual_id] = molecule_node

        for individual in minerva_obj["individuals"]:
            individual_id = individual["id"]
            root_types = [x["id"] for x in individual.get("root-type", []) if x]
            individual_to_root_types[individual_id] = root_types

            term_id: str | None = None
            for type_ in individual.get("type", []):
                if type_.get("type") == "complement":
                    # class expression representing NOT
                    continue
                type_id = type_.get("id")
                if type_id is None:
                    continue
                objects_by_id[type_id] = type_
                term_id = type_id

            if term_id is None:
                translation_warnings.add(
                    TranslationWarning(
                        type=WarningType.MISSING_TERM,
                        message=f"Missing term for individual {individual_id}",
                        entity_id=individual_id,
                    )
                )
                continue

            individual_to_term[individual_id] = term_id
            individual_to_annotations[individual_id] = _annotations(individual)
            individual_to_annotations_multivalued[individual_id] = (
                _annotations_multivalued(individual)
            )

        for fact in minerva_obj["facts"]:
            facts_by_property[fact["property"]].append(fact)
            facts_by_subject_property[(fact["subject"], fact["property"])].append(fact)

        enabled_by_facts = facts_by_property.get(ENABLED_BY, [])
        if not enabled_by_facts:
            translation_warnings.add(
                TranslationWarning(
                    type=WarningType.NO_ENABLED_BY_FACTS,
                    message=f"No enabled_by ({ENABLED_BY}) facts found",
                    entity_id=id,
                )
            )
        for enabled_by_fact in enabled_by_facts:
            subject, object_ = enabled_by_fact["subject"], enabled_by_fact["object"]
            if subject not in individual_to_term:
                translation_warnings.add(
                    TranslationWarning(
                        type=WarningType.MISSING_TERM,
                        message=f"Missing term for subject {subject} in fact",
                        entity_id=subject,
                    )
                )
                continue
            if object_ not in individual_to_term:
                translation_warnings.add(
                    TranslationWarning(
                        type=WarningType.MISSING_TERM,
                        message=f"Missing term for object {object_} in fact",
                        entity_id=object_,
                    )
                )
                continue
            gene_id = individual_to_term[object_]
            root_types = individual_to_root_types.get(object_, [])

            evs, prov = _process_fact(enabled_by_fact)
            enabled_by_association: EnabledByAssociation
            if PROTEIN_CONTAINING_COMPLEX in root_types:
                member_associations: list[ProteinComplexMemberAssociation] = []
                has_part_facts = facts_by_subject_property.get((object_, HAS_PART), [])
                for has_part_fact in has_part_facts:
                    member_object = has_part_fact["object"]
                    member_term = individual_to_term.get(member_object)
                    if member_term is None:
                        translation_warnings.add(
                            TranslationWarning(
                                type=WarningType.MISSING_TERM,
                                message=f"Missing term for object {member_object} in has_part fact",
                                entity_id=member_object,
                            )
                        )
                        continue
                    member_evs, member_prov = _process_fact(has_part_fact)
                    member_associations.append(
                        ProteinComplexMemberAssociation(
                            term=member_term,
                            evidence=member_evs,
                            provenances=[member_prov],
                        )
                    )
                enabled_by_association = EnabledByProteinComplexAssociation(
                    term=gene_id,
                    members=member_associations,
                    evidence=evs,
                    provenances=[prov],
                )
            elif INFORMATION_BIOMACROMOLECULE in root_types:
                enabled_by_association = EnabledByGeneProductAssociation(
                    term=gene_id, evidence=evs, provenances=[prov]
                )
            else:
                translation_warnings.add(
                    TranslationWarning(
                        type=WarningType.UNKNOWN_ENABLED_BY_TYPE,
                        message=f"Unknown enabled_by type for {object_}; assuming gene product",
                        entity_id=object_,
                    )
                )
                enabled_by_association = EnabledByGeneProductAssociation(
                    term=gene_id, evidence=evs, provenances=[prov]
                )

            activity = Activity(
                id=subject,
                enabled_by=enabled_by_association,
                molecular_function=MolecularFunctionAssociation(
                    term=individual_to_term[subject]
                ),
            )
            activities.append(activity)
            activities_by_mf_id[subject].append(activity)

        for activity, object_, evs, prov in _iter_activities_by_fact_subject(
            fact_property=PART_OF
        ):
            association = BiologicalProcessAssociation(
                term=individual_to_term[object_], evidence=evs, provenances=[prov]
            )

            happens_during_facts = facts_by_subject_property.get(
                (object_, HAPPENS_DURING), []
            )
            for happens_during_fact in happens_during_facts:
                evs, prov = _process_fact(happens_during_fact)
                if association.happens_during is not None:
                    translation_warnings.add(
                        TranslationWarning(
                            type=WarningType.ATTRIBUTE_OVERWRITE,
                            message=f"Overwriting part_of.happens_during for {activity.id}",
                            entity_id=activity.id,
                        )
                    )
                association.happens_during = PhaseAssociation(
                    term=individual_to_term.get(happens_during_fact["object"], None),
                    evidence=evs,
                    provenances=[prov],
                )

            part_of_facts = facts_by_subject_property.get((object_, PART_OF), [])
            for part_of_fact in part_of_facts:
                evs, prov = _process_fact(part_of_fact)
                if association.part_of is not None:
                    translation_warnings.add(
                        TranslationWarning(
                            type=WarningType.ATTRIBUTE_OVERWRITE,
                            message=f"Overwriting part_of.part_of for {activity.id}",
                            entity_id=activity.id,
                        )
                    )
                association.part_of = BiologicalProcessAssociation(
                    term=individual_to_term.get(part_of_fact["object"], None),
                    evidence=evs,
                    provenances=[prov],
                )

            _setattr_with_warning(activity, "part_of", association)

        for activity, object_, evs, prov in _iter_activities_by_fact_subject(
            fact_property=OCCURS_IN
        ):
            association = CellularAnatomicalEntityAssociation(
                term=individual_to_term[object_], evidence=evs, provenances=[prov]
            )

            part_of_facts = facts_by_subject_property.get((object_, PART_OF), [])
            for part_of_fact in part_of_facts:
                evs, prov = _process_fact(part_of_fact)
                if association.part_of is not None:
                    translation_warnings.add(
                        TranslationWarning(
                            type=WarningType.ATTRIBUTE_OVERWRITE,
                            message=f"Overwriting occurs_in.part_of for {activity.id}",
                            entity_id=activity.id,
                        )
                    )
                association.part_of = CellTypeAssociation(
                    term=individual_to_term.get(part_of_fact["object"], None),
                    evidence=evs,
                    provenances=[prov],
                )

            _setattr_with_warning(activity, "occurs_in", association)

        for activity, object_, evs, prov in _iter_activities_by_fact_subject(
            fact_property=HAPPENS_DURING
        ):
            association = PhaseAssociation(
                term=individual_to_term[object_], evidence=evs, provenances=[prov]
            )
            _setattr_with_warning(activity, "happens_during", association)

        for activity, object_, evs, prov in _iter_activities_by_fact_subject(
            fact_property=HAS_INPUT
        ):
            if activity.has_input is None:
                activity.has_input = []
            _add_molecule_node(object_)
            activity.has_input.append(
                MoleculeAssociation(molecule=object_, evidence=evs, provenances=[prov])
            )

        for activity, object_, evs, prov in _iter_activities_by_fact_subject(
            fact_property=HAS_PRIMARY_INPUT
        ):
            _add_molecule_node(object_)
            association = MoleculeAssociation(
                molecule=object_, evidence=evs, provenances=[prov]
            )
            _setattr_with_warning(activity, "has_primary_input", association)

        for activity, object_, evs, prov in _iter_activities_by_fact_subject(
            fact_property=HAS_OUTPUT
        ):
            if activity.has_output is None:
                activity.has_output = []
            _add_molecule_node(object_)
            activity.has_output.append(
                MoleculeAssociation(molecule=object_, evidence=evs, provenances=[prov])
            )

        for activity, object_, evs, prov in _iter_activities_by_fact_subject(
            fact_property=HAS_PRIMARY_OUTPUT
        ):
            _add_molecule_node(object_)
            association = MoleculeAssociation(
                molecule=object_, evidence=evs, provenances=[prov]
            )
            _setattr_with_warning(activity, "has_primary_output", association)

        for fact_property, facts in facts_by_property.items():
            for fact in facts:
                subject, object_ = fact["subject"], fact["object"]
                subject_activities = activities_by_mf_id.get(subject, [])
                object_activities = activities_by_mf_id.get(object_, [])

                if not subject_activities or not object_activities:
                    continue
                if MOLECULAR_FUNCTION not in individual_to_root_types.get(subject, []):
                    continue
                if MOLECULAR_FUNCTION not in individual_to_root_types.get(object_, []):
                    continue
                if len(subject_activities) > 1:
                    translation_warnings.add(
                        TranslationWarning(
                            type=WarningType.MULTIPLE_ACTIVITIES,
                            message=f"Multiple activities for subject {subject}",
                            entity_id=subject,
                        )
                    )
                if len(object_activities) > 1:
                    translation_warnings.add(
                        TranslationWarning(
                            type=WarningType.MULTIPLE_ACTIVITIES,
                            message=f"Multiple activities for object {object_}",
                            entity_id=object_,
                        )
                    )

                subject_activity = subject_activities[0]
                object_activity = object_activities[0]
                evs, provenance = _process_fact(fact)
                rel = CausalAssociation(
                    predicate=fact_property,
                    downstream_activity=object_activity.id,
                    evidence=evs,
                    provenances=[provenance],
                )
                if subject_activity.causal_associations is None:
                    subject_activity.causal_associations = []
                subject_activity.causal_associations.append(rel)

        annotations = _annotations(minerva_obj)
        annotations_mv = _annotations_multivalued(minerva_obj)

        objects: list[Object] = []
        for obj in objects_by_id.values():
            object_ = Object(id=obj["id"])
            if "label" in obj:
                object_.label = obj["label"]
            objects.append(object_)

        provenance = ProvenanceInfo(
            contributor=annotations_mv.get("contributor"),
            date=annotations.get("date", None),
            provided_by=annotations_mv.get("providedBy"),
        )

        # Get all taxa from the annotations
        all_taxa = annotations_mv.get(TaxonVocabulary.TAXON_ANNOTATION_KEY, [])

        # Add legacy taxon key if it exists (backward compatibility)
        legacy_taxon = annotations.get(TaxonVocabulary.LEGACY_TAXON_KEY)
        if legacy_taxon and legacy_taxon not in all_taxa:
            all_taxa.append(legacy_taxon)

        # If no taxa, nothing to do
        if not all_taxa:
            taxon = None
            additional_taxa = []
        # If only one taxon, it's the primary
        elif len(all_taxa) == 1:
            taxon = all_taxa[0]
            additional_taxa = []
        # Multiple taxa - prioritize host taxa as primary
        else:
            # Find host taxa in the list
            host_matches = [t for t in all_taxa if TaxonVocabulary.is_host_taxon(t)]
            if host_matches:
                # Use the first host taxon as primary
                taxon = host_matches[0]
                # All others are additional taxa
                additional_taxa = [t for t in all_taxa if t != taxon]
            else:
                # No host matches, just use the first as primary
                taxon = all_taxa[0]
                additional_taxa = all_taxa[1:]

        num_processed_facts = len(processed_facts)
        num_facts = len(minerva_obj["facts"])
        if num_processed_facts < num_facts:
            for fact in minerva_obj["facts"]:
                fact_tuple = (fact["subject"], fact["property"], fact["object"])
                if fact_tuple not in processed_facts:
                    translation_warnings.add(
                        TranslationWarning(
                            type=WarningType.UNHANDLED_FACT,
                            message=f"Unhandled fact: {fact_tuple}",
                            entity_id=id,
                        )
                    )
        cam = Model(
            id=id,
            title=annotations["title"],
            comments=annotations_mv.get("comment", None),
            date_modified=annotations.get("date", None),
            taxon=taxon,
            activities=activities,
            molecules=list(molecule_nodes_by_id.values()),
            objects=objects,
            provenances=[provenance],
        )

        state_annotation = annotations.get("state", None)
        if state_annotation:
            try:
                cam.status = ModelStateEnum(state_annotation)
            except ValueError:
                translation_warnings.add(
                    TranslationWarning(
                        type=WarningType.INVALID_MODEL_STATE,
                        message=f"Invalid model state {state_annotation}",
                        entity_id=id,
                    )
                )

        if additional_taxa and len(additional_taxa) > 0:
            cam.additional_taxa = additional_taxa

        return TranslationResult(result=cam, warnings=list(translation_warnings))

    @staticmethod
    def minerva_object_to_model(obj: dict) -> Model:
        """Convert a Minerva JSON object to a GO-CAM Model.

        Args:
            obj: Minerva JSON object

        Returns:
            GO-CAM Model
        """
        result = MinervaWrapper.translate(obj)
        return result.result
