#
#   Imandra Inc.
#
#   speclogician/data/container.py
#

from typing import Any
from pydantic import BaseModel, Field

from .traces import TestTrace, LogTrace
from .refs import DocRef, SrcCodeRef
from .mapping import ArtifactMap

from ..utils import IMLValidity

class SmartFindResults(BaseModel):
    """Helper for returning search results"""
    needle: str

    test_traces: list[TestTrace] = Field(default_factory=list[TestTrace])
    log_traces: list[LogTrace] = Field(default_factory=list[LogTrace])
    doc_refs: list[DocRef] = Field(default_factory=list[DocRef])
    src_code_refs: list[SrcCodeRef] = Field(default_factory=list[SrcCodeRef])

    def empty(self) -> bool:
        return not (self.test_traces or self.log_traces or self.doc_refs or self.src_code_refs)


class ArtifactContainer(BaseModel):
    """Artifact container"""

    test_traces: list[TestTrace] = Field(default_factory=list[TestTrace])
    log_traces: list[LogTrace] = Field(default_factory=list[LogTrace])
    doc_ref: list[DocRef] = Field(default_factory=list[DocRef])
    src_code: list[SrcCodeRef] = Field(default_factory=list[SrcCodeRef])

    num_test_traces_total: int = 0
    num_test_traces_matched: int = 0
    num_test_traces_logic_good: int = 0

    num_log_traces_total: int = 0
    num_log_traces_matched: int = 0
    num_log_traces_logic_good: int = 0

    num_src_code_arts_matched: int = 0
    num_src_code_arts_total: int = 0

    num_doc_arts_matched: int = 0
    num_doc_arts_total: int = 0

    num_test_trace_steps_total: int = 0
    num_test_trace_steps_matched: int = 0

    num_log_trace_steps_total: int = 0
    num_log_trace_steps_matched: int = 0

    # Optional: how many traces have at least one matched step
    num_test_traces_with_any_step_matched: int = 0
    num_log_traces_with_any_step_matched: int = 0


    # ----------------------------
    # Internals
    # ----------------------------

    def _all_art_ids(self) -> set[str]:
        ids: set[str] = set()
        for xs in (self.test_traces, self.log_traces, self.doc_ref, self.src_code):
            for a in xs:
                if a.art_id is not None:
                    ids.add(a.art_id)
        return ids

    def _ensure_unique_art_id(self, art_id: str) -> None:
        if art_id in self._all_art_ids():
            raise ValueError(f"Artifact with art_id='{art_id}' already exists in container")

    def refresh_stats(self, art_map: ArtifactMap | None = None) -> None:
        # Totals (artifact counts)
        self.num_test_traces_total = len(self.test_traces)
        self.num_log_traces_total = len(self.log_traces)
        self.num_src_code_arts_total = len(self.src_code)
        self.num_doc_arts_total = len(self.doc_ref)

        # Step totals
        self.num_test_trace_steps_total = sum(len(t.steps) for t in self.test_traces)
        self.num_log_trace_steps_total = sum(len(t.steps) for t in self.log_traces)

        # Matched (trace-level + step-level)
        if art_map is not None:
            def _is_linked(art_id: str | None) -> bool:
                if not art_id:
                    return False
                return len(art_map.get_comp_for_artifact(art_id)) > 0

            def _step_ref(trace_id: str, step_idx: int) -> str:
                return f"trace:{trace_id}:step:{step_idx}"

            # trace-level matched (by art_id)
            self.num_test_traces_matched = sum(1 for t in self.test_traces if _is_linked(t.art_id))
            self.num_log_traces_matched = sum(1 for t in self.log_traces if _is_linked(t.art_id))
            self.num_src_code_arts_matched = sum(1 for a in self.src_code if _is_linked(a.art_id))
            self.num_doc_arts_matched = sum(1 for a in self.doc_ref if _is_linked(a.art_id))

            # step-level matched
            test_steps_matched = 0
            test_traces_any_step = 0
            for t in self.test_traces:
                if not t.art_id:
                    continue
                any_step = False
                for i in range(len(t.steps)):
                    if _is_linked(_step_ref(t.art_id, i)):
                        test_steps_matched += 1
                        any_step = True
                if any_step:
                    test_traces_any_step += 1

            log_steps_matched = 0
            log_traces_any_step = 0
            for t in self.log_traces:
                if not t.art_id:
                    continue
                any_step = False
                for i in range(len(t.steps)):
                    if _is_linked(_step_ref(t.art_id, i)):
                        log_steps_matched += 1
                        any_step = True
                if any_step:
                    log_traces_any_step += 1

            self.num_test_trace_steps_matched = test_steps_matched
            self.num_log_trace_steps_matched = log_steps_matched
            self.num_test_traces_with_any_step_matched = test_traces_any_step
            self.num_log_traces_with_any_step_matched = log_traces_any_step

        else:
            # keep existing behavior: if no map, "matched" stays 0
            self.num_test_traces_matched = 0
            self.num_log_traces_matched = 0
            self.num_src_code_arts_matched = 0
            self.num_doc_arts_matched = 0

            self.num_test_trace_steps_matched = 0
            self.num_log_trace_steps_matched = 0
            self.num_test_traces_with_any_step_matched = 0
            self.num_log_traces_with_any_step_matched = 0

        # Logic good: count VALID IML traces (unchanged)
        self.num_test_traces_logic_good = sum(
            1
            for t in self.test_traces
            if getattr(getattr(t, "iml_validity", None), "iml_valid", IMLValidity.UNKNOWN) == IMLValidity.VALID
        )
        self.num_log_traces_logic_good = sum(
            1
            for t in self.log_traces
            if getattr(getattr(t, "iml_validity", None), "iml_valid", IMLValidity.UNKNOWN) == IMLValidity.VALID
        )
    # ----------------------------
    # Traces
    # ----------------------------

    def add_trace(self, trace: TestTrace | LogTrace) -> None:
        """Add a trace (TestTrace or LogTrace). Uniqueness is by art_id only."""
        if trace.art_id is None:
            raise ValueError("Provided art_id is None")
        self._ensure_unique_art_id(trace.art_id)

        if isinstance(trace, TestTrace):
            self.test_traces.append(trace)
        elif isinstance(trace, LogTrace):
            self.log_traces.append(trace)
        else:
            raise TypeError(f"Unsupported trace type: {type(trace).__name__}")

        self.refresh_stats()

    def rem_trace(self, art_id: str) -> None:
        """Remove exactly one trace by art_id."""
        before = len(self.test_traces)
        self.test_traces = [t for t in self.test_traces if t.art_id != art_id]
        if len(self.test_traces) != before:
            self.refresh_stats()
            return

        before = len(self.log_traces)
        self.log_traces = [t for t in self.log_traces if t.art_id != art_id]
        if len(self.log_traces) != before:
            self.refresh_stats()
            return

        raise ValueError(f"Trace not found (art_id='{art_id}')")

    def edit_trace(self, art_id: str, patch: dict[str, object]) -> None:
        """
        Patch-edit a trace by art_id (Pydantic model_copy(update=...)).
        """
        for i, t in enumerate(self.test_traces):
            if t.art_id == art_id:
                self.test_traces[i] = t.model_copy(update=patch)
                self.refresh_stats()
                return

        for i, t in enumerate(self.log_traces):
            if t.art_id == art_id:
                self.log_traces[i] = t.model_copy(update=patch)
                self.refresh_stats()
                return

        raise ValueError(f"Trace not found (art_id='{art_id}')")

    # ----------------------------
    # Data artifacts
    # ----------------------------

    def add_data_art(self, art: DocRef | SrcCodeRef) -> None:
        """Add a data artifact (DocRef or SrcCodeRef)."""
        if art.art_id is None:
            raise ValueError("Provided art_id is None")
        self._ensure_unique_art_id(art.art_id)

        if isinstance(art, DocRef):
            self.doc_ref.append(art)
        elif isinstance(art, SrcCodeRef):
            self.src_code.append(art)
        else:
            raise TypeError(f"Unsupported data artifact type: {type(art).__name__}")

        self.refresh_stats()

    def rem_data_art(self, art_id: str) -> None:
        """Remove exactly one data artifact by art_id."""
        before = len(self.doc_ref)
        self.doc_ref = [a for a in self.doc_ref if a.art_id != art_id]
        if len(self.doc_ref) != before:
            self.refresh_stats()
            return

        before = len(self.src_code)
        self.src_code = [a for a in self.src_code if a.art_id != art_id]
        if len(self.src_code) != before:
            self.refresh_stats()
            return

        raise ValueError(f"Data artifact not found (art_id='{art_id}')")

    def edit_data_art(self, art_id: str, patch: dict[str, object]) -> None:
        """Patch-edit a data artifact by art_id."""
        for i, a in enumerate(self.doc_ref):
            if a.art_id == art_id:
                self.doc_ref[i] = a.model_copy(update=patch)
                self.refresh_stats()
                return

        for i, a in enumerate(self.src_code):
            if a.art_id == art_id:
                self.src_code[i] = a.model_copy(update=patch)
                self.refresh_stats()
                return

        raise ValueError(f"Data artifact not found (art_id='{art_id}')")

    # ----------------------------
    # Iterators / basic accessors
    # ----------------------------

    def all_traces(self) -> list[TestTrace | LogTrace]:
        return list(self.test_traces) + list(self.log_traces)

    def all_data_arts(self) -> list[DocRef | SrcCodeRef]:
        return list(self.doc_ref) + list(self.src_code)

    def get_by_art_id(self, art_id: str) -> TestTrace | LogTrace | DocRef | SrcCodeRef | None:
        for a in self.test_traces:
            if a.art_id == art_id:
                return a
        for a in self.log_traces:
            if a.art_id == art_id:
                return a
        for a in self.doc_ref:
            if a.art_id == art_id:
                return a
        for a in self.src_code:
            if a.art_id == art_id:
                return a
        return None


    # ----------------------------
    # TestTrace lookups
    # ----------------------------

    def find_test_traces_by_name(self, name: str) -> list[TestTrace]:
        # NOT unique by design
        return [t for t in self.test_traces if t.name == name]

    def find_test_traces_by_filepath(self, filepath: str) -> list[TestTrace]:
        return [t for t in self.test_traces if t.filepath == filepath]

    def find_test_traces_by_language(self, language: str) -> list[TestTrace]:
        return [t for t in self.test_traces if (t.language or "") == language]

    def search_test_traces_text(self, needle: str) -> list[TestTrace]:
        n = needle.lower()
        out: list[TestTrace] = []
        for t in self.test_traces:
            hay = " ".join(
                [
                    t.name or "",
                    t.filepath or "",
                    t.language or "",
                    t.contents or "",
                ]
            ).lower()
            if n in hay:
                out.append(t)
        return out


    # ----------------------------
    # LogTrace lookups
    # ----------------------------

    def find_log_traces_by_filename(self, filename: str) -> list[LogTrace]:
        # NOT unique by design
        return [t for t in self.log_traces if t.filename == filename]

    def search_log_traces_text(self, needle: str) -> list[LogTrace]:
        n = needle.lower()
        out: list[LogTrace] = []
        for t in self.log_traces:
            hay = " ".join([t.filename or "", t.contents or ""]).lower()
            if n in hay:
                out.append(t)
        return out


    # ----------------------------
    # Data artifact lookups
    # ----------------------------

    def find_doc_refs_containing(self, needle: str) -> list[DocRef]:
        n = needle.lower()
        return [
            d
            for d in self.doc_ref
            if n in (d.text or "").lower() or n in ((d.meta or "").lower())
        ]

    def find_src_code_by_file_path(self, file_path: str) -> list[SrcCodeRef]:
        return [s for s in self.src_code if (s.file_path or "") == file_path]

    def find_src_code_by_language(self, language: str) -> list[SrcCodeRef]:
        return [s for s in self.src_code if (s.language or "") == language]

    def search_src_code_text(self, needle: str) -> list[SrcCodeRef]:
        n = needle.lower()
        out: list[SrcCodeRef] = []
        for s in self.src_code:
            hay = " ".join(
                [
                    s.file_path or "",
                    s.language or "",
                    s.src_code or "",
                    s.iml_code or "",
                    s.meta or "",
                ]
            ).lower()
            if n in hay:
                out.append(s)
        return out

    def smart_find(
        self,
        needle: str,
        *,
        include: set[str] | None = None,
        max_per_type: int | None = 50,
    ) -> SmartFindResults:
        """
        Search across *all* artifacts with a single entrypoint.

        - `include`: optional subset of {"test", "log", "doc", "src"}; default = all.
        - `max_per_type`: cap results per type (None = unlimited).
        """
        inc = include or {"test", "log", "doc", "src"}
        n = (needle or "").lower()

        def _cap(xs: list[Any]) -> list[Any]:
            if max_per_type is None:
                return xs
            return xs[:max_per_type]

        out = SmartFindResults(needle=needle)

        if "test" in inc:
            hits: list[TestTrace] = []
            for t in self.test_traces:
                hay = " ".join(
                    [
                        t.art_id or "",
                        t.name or "",
                        t.filepath or "",
                        t.language or "",
                        t.contents or "",
                        t.given or "",
                        t.when or "",
                        t.then or "",
                    ]
                ).lower()
                if n in hay:
                    hits.append(t)
            out.test_traces = _cap(hits)

        if "log" in inc:
            hits: list[LogTrace] = []
            for t in self.log_traces:
                hay = " ".join(
                    [
                        t.art_id or "",
                        t.filename or "",
                        t.contents or "",
                        t.given or "",
                        t.when or "",
                        t.then or "",
                    ]
                ).lower()
                if n in hay:
                    hits.append(t)
            out.log_traces = _cap(hits)

        if "doc" in inc:
            hits: list[DocRef] = []
            for d in self.doc_ref:
                hay = " ".join(
                    [
                        d.art_id or "",
                        d.text or "",
                        d.meta or "",
                    ]
                ).lower()
                if n in hay:
                    hits.append(d)
            out.doc_refs = _cap(hits)

        if "src" in inc:
            hits: list[SrcCodeRef] = []
            for s in self.src_code:
                hay = " ".join(
                    [
                        s.art_id or "",
                        s.file_path or "",
                        s.language or "",
                        s.src_code or "",
                        s.iml_code or "",
                        s.meta or "",
                    ]
                ).lower()
                if n in hay:
                    hits.append(s)
            out.src_code_refs = _cap(hits)

        return out
