"""Tenant-scoped SQLAlchemy repository for local graph memory."""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Any

from sqlalchemy import or_, select
from sqlalchemy.orm import Session

from ..models.memory import MemoryEdge, MemoryEpisode, MemoryGraph, MemoryNode


@dataclass(frozen=True)
class MemorySearchResult:
    nodes: list[MemoryNode] = field(default_factory=list)
    edges: list[MemoryEdge] = field(default_factory=list)
    query: str = ""

    @property
    def total_count(self) -> int:
        return len(self.nodes) + len(self.edges)


class SqlAlchemyMemoryRepository:
    """All reads and writes are constrained by organization_id + graph_id."""

    def __init__(self, session: Session, *, organization_id: str, graph_id: str):
        if not isinstance(organization_id, str) or not organization_id.strip():
            raise ValueError("invalid_organization_id")
        if not isinstance(graph_id, str) or not graph_id.strip():
            raise ValueError("invalid_graph_id")
        self.session = session
        self.organization_id = organization_id
        self.graph_id = graph_id

    def _graph(self) -> MemoryGraph | None:
        return self.session.scalar(
            select(MemoryGraph).where(
                MemoryGraph.id == self.graph_id,
                MemoryGraph.organization_id == self.organization_id,
            )
        )

    def _require_graph(self) -> MemoryGraph:
        graph = self._graph()
        if graph is None:
            raise ValueError("memory_graph_not_found")
        return graph

    def get_graph(self) -> MemoryGraph:
        """Return the graph only when it belongs to this organization scope."""
        return self._require_graph()

    def update_graph_ontology(self, ontology: dict[str, Any]) -> MemoryGraph:
        if not isinstance(ontology, dict):
            raise ValueError("invalid_ontology")
        graph = self._require_graph()
        graph.ontology = dict(ontology)
        graph.version = int(graph.version or 0) + 1
        self.session.flush()
        return graph

    def get_episode(self, *, source_type: str, source_ref: str) -> MemoryEpisode | None:
        self._require_graph()
        return self.session.scalar(
            select(MemoryEpisode).where(
                MemoryEpisode.graph_id == self.graph_id,
                MemoryEpisode.source_type == source_type,
                MemoryEpisode.source_ref == source_ref,
            )
        )

    @staticmethod
    def _normalize_name(name: str) -> str:
        if not isinstance(name, str) or not name.strip():
            raise ValueError("invalid_memory_name")
        return " ".join(name.casefold().split())

    @staticmethod
    def _confidence(value: float) -> float:
        if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
            raise ValueError("invalid_confidence")
        return min(max(float(value), 0.0), 1.0)

    def create_graph(self, *, project_id: str, ontology: dict[str, Any] | None = None) -> MemoryGraph:
        if not isinstance(project_id, str) or not project_id.strip():
            raise ValueError("invalid_project_id")
        existing = self.session.get(MemoryGraph, self.graph_id)
        if existing is not None:
            if existing.organization_id != self.organization_id:
                raise ValueError("memory_graph_scope_conflict")
            return existing
        graph = MemoryGraph(
            id=self.graph_id,
            organization_id=self.organization_id,
            project_id=project_id,
            ontology=ontology or {},
        )
        self.session.add(graph)
        self.session.flush()
        return graph

    def add_episode(
        self,
        *,
        source_type: str,
        source_ref: str,
        normalized_text: str,
        summary: str = "",
        extractor_version: str = "v1",
    ) -> MemoryEpisode:
        self._require_graph()
        if not all(isinstance(value, str) and value.strip() for value in (source_type, source_ref, normalized_text)):
            raise ValueError("invalid_memory_episode")
        episode = self.session.scalar(
            select(MemoryEpisode).where(
                MemoryEpisode.graph_id == self.graph_id,
                MemoryEpisode.source_type == source_type,
                MemoryEpisode.source_ref == source_ref,
            )
        )
        if episode is None:
            episode = MemoryEpisode(
                graph_id=self.graph_id,
                source_type=source_type,
                source_ref=source_ref,
                normalized_text=normalized_text,
                summary=summary or "",
                extractor_version=extractor_version or "v1",
            )
            self.session.add(episode)
        else:
            episode.normalized_text = normalized_text
            episode.summary = summary or ""
            episode.extractor_version = extractor_version or "v1"
        self.session.flush()
        return episode

    def upsert_node(
        self,
        *,
        canonical_name: str,
        labels: list[str] | None = None,
        aliases: list[str] | None = None,
        attributes: dict[str, Any] | None = None,
        summary: str = "",
        confidence: float = 0.0,
    ) -> MemoryNode:
        self._require_graph()
        normalized_name = self._normalize_name(canonical_name)
        node = self.session.scalar(
            select(MemoryNode).where(
                MemoryNode.graph_id == self.graph_id,
                MemoryNode.normalized_name == normalized_name,
            )
        )
        if node is None:
            node = MemoryNode(graph_id=self.graph_id, canonical_name=canonical_name.strip(), normalized_name=normalized_name)
            self.session.add(node)
        node.canonical_name = canonical_name.strip()
        node.labels = list(labels or [])
        node.aliases = list(aliases or [])
        node.attributes = dict(attributes or {})
        node.summary = summary or ""
        node.confidence = self._confidence(confidence)
        self.session.flush()
        return node

    def upsert_edge(
        self,
        *,
        source_node_id: str,
        target_node_id: str,
        relation: str,
        fact: str,
        attributes: dict[str, Any] | None = None,
        confidence: float = 0.0,
        valid_at=None,
        invalid_at=None,
        expired_at=None,
    ) -> MemoryEdge:
        self._require_graph()
        if not all(isinstance(value, str) and value.strip() for value in (source_node_id, target_node_id, relation, fact)):
            raise ValueError("invalid_memory_edge")
        nodes = list(
            self.session.scalars(
                select(MemoryNode).where(
                    MemoryNode.graph_id == self.graph_id,
                    MemoryNode.id.in_([source_node_id, target_node_id]),
                )
            )
        )
        if {node.id for node in nodes} != {source_node_id, target_node_id}:
            raise ValueError("memory_edge_node_scope_conflict")
        edge = self.session.scalar(
            select(MemoryEdge).where(
                MemoryEdge.graph_id == self.graph_id,
                MemoryEdge.source_node_id == source_node_id,
                MemoryEdge.target_node_id == target_node_id,
                MemoryEdge.relation == relation.strip(),
                MemoryEdge.fact == fact.strip(),
            )
        )
        if edge is None:
            edge = MemoryEdge(
                graph_id=self.graph_id,
                source_node_id=source_node_id,
                target_node_id=target_node_id,
                relation=relation.strip(),
                fact=fact.strip(),
            )
            self.session.add(edge)
        edge.attributes = dict(attributes or {})
        edge.confidence = self._confidence(confidence)
        edge.valid_at = valid_at
        edge.invalid_at = invalid_at
        edge.expired_at = expired_at
        self.session.flush()
        return edge

    def get_node(self, node_id: str) -> MemoryNode | None:
        return self.session.scalar(
            select(MemoryNode)
            .join(MemoryGraph, MemoryGraph.id == MemoryNode.graph_id)
            .where(
                MemoryNode.id == node_id,
                MemoryNode.graph_id == self.graph_id,
                MemoryGraph.organization_id == self.organization_id,
            )
        )

    def list_nodes(self, *, limit: int = 1_000) -> list[MemoryNode]:
        self._require_graph()
        safe_limit = min(max(int(limit), 1), 10_000)
        return list(
            self.session.scalars(
                select(MemoryNode)
                .join(MemoryGraph, MemoryGraph.id == MemoryNode.graph_id)
                .where(
                    MemoryNode.graph_id == self.graph_id,
                    MemoryGraph.organization_id == self.organization_id,
                )
                .order_by(MemoryNode.normalized_name.asc(), MemoryNode.id.asc())
                .limit(safe_limit)
            )
        )

    def list_edges(self, *, limit: int = 2_000) -> list[MemoryEdge]:
        self._require_graph()
        safe_limit = min(max(int(limit), 1), 20_000)
        return list(
            self.session.scalars(
                select(MemoryEdge)
                .join(MemoryGraph, MemoryGraph.id == MemoryEdge.graph_id)
                .where(
                    MemoryEdge.graph_id == self.graph_id,
                    MemoryGraph.organization_id == self.organization_id,
                )
                .order_by(MemoryEdge.id.asc())
                .limit(safe_limit)
            )
        )

    def get_node_edges(self, node_id: str, *, limit: int = 500) -> list[MemoryEdge]:
        self._require_graph()
        safe_limit = min(max(int(limit), 1), 5_000)
        return list(
            self.session.scalars(
                select(MemoryEdge)
                .join(MemoryGraph, MemoryGraph.id == MemoryEdge.graph_id)
                .where(
                    MemoryEdge.graph_id == self.graph_id,
                    MemoryGraph.organization_id == self.organization_id,
                    (MemoryEdge.source_node_id == node_id) | (MemoryEdge.target_node_id == node_id),
                )
                .order_by(MemoryEdge.id.asc())
                .limit(safe_limit)
            )
        )

    def get_node_by_id(self, node_id: str) -> MemoryNode | None:
        return self.get_node(node_id)

    def search(self, query: str, *, limit: int = 50) -> MemorySearchResult:
        if not isinstance(query, str):
            raise ValueError("invalid_memory_query")
        normalized_query = " ".join(query.casefold().split())
        if not normalized_query:
            return MemorySearchResult(query=query)
        safe_limit = min(max(int(limit), 1), 100)
        pattern = f"%{normalized_query}%"
        nodes = list(
            self.session.scalars(
                select(MemoryNode)
                .join(MemoryGraph, MemoryGraph.id == MemoryNode.graph_id)
                .where(
                    MemoryNode.graph_id == self.graph_id,
                    MemoryGraph.organization_id == self.organization_id,
                    or_(
                        MemoryNode.normalized_name.ilike(pattern),
                        MemoryNode.summary.ilike(pattern),
                    ),
                )
                .order_by(MemoryNode.normalized_name.asc(), MemoryNode.id.asc())
                .limit(safe_limit)
            )
        )
        return MemorySearchResult(nodes=nodes, query=query)