"""Repository methods for docsctx."""

from __future__ import annotations

import json
import hashlib
from dataclasses import asdict, dataclass
from pathlib import Path

from sqlalchemy import func, select
from sqlalchemy.orm import selectinload

from common.db import Base, create_engine_with_wal, create_session_factory, session_scope
from docsctx.analysis import CommunityResult
from docsctx.db.models import Document, DocumentChunk, DocumentCommunity, DocumentCommunityMembership, DocumentEntity, DocumentRelationship, DocumentClaim


@dataclass(slots=True)
class DocumentRecord:
    source_uri: str
    title: str
    doc_type: str
    checksum: str
    content: str
    chunks: list[str]


@dataclass(slots=True)
class DocumentStats:
    documents: int = 0
    chunks: int = 0
    embeddings: int = 0
    entities: int = 0
    relationships: int = 0
    claims: int = 0
    communities: int = 0

    def to_dict(self) -> dict[str, int]:
        return asdict(self)


@dataclass(slots=True)
class SearchResult:
    document_id: int
    source_uri: str
    title: str
    doc_type: str
    chunk_id: int
    chunk_index: int
    snippet: str
    score: float
    is_latest: bool = True

    def to_dict(self) -> dict[str, object]:
        return asdict(self)


@dataclass(slots=True)
class CommunitySummary:
    id: int
    label: str
    summary: str
    top_terms: list[str]
    document_ids: list[int]

    def to_dict(self) -> dict[str, object]:
        return asdict(self)


class DocsRepository:
    """Persistence wrapper for docsctx."""

    def __init__(self, db_path: str | Path, echo: bool = False) -> None:
        self.db_path = Path(db_path)
        self.engine = create_engine_with_wal(f"sqlite:///{self.db_path}", echo=echo)
        self.session_factory = create_session_factory(self.engine)

    def create_schema(self) -> None:
        Base.metadata.create_all(
            self.engine,
            tables=[
                Document.__table__,
                DocumentChunk.__table__,
                DocumentEntity.__table__,
                DocumentRelationship.__table__,
                DocumentClaim.__table__,
                DocumentCommunity.__table__,
                DocumentCommunityMembership.__table__,
            ],
        )

    def clear(self) -> None:
        """Remove all docsctx data and recreate the schema."""
        Base.metadata.drop_all(self.engine)
        self.create_schema()

    def upsert_document(self, record: DocumentRecord) -> int:
        with session_scope(self.session_factory) as session:
            existing = session.scalar(
                select(Document)
                .options(
                    selectinload(Document.chunks),
                    selectinload(Document.entities),
                    selectinload(Document.relationships),
                )
                .where(Document.source_uri == record.source_uri)
                .where(Document.is_latest == True)
            )
            if existing is None:
                canonical_id = hashlib.sha256(record.source_uri.encode()).hexdigest()[:16]
                document = Document(
                    source_uri=record.source_uri,
                    title=record.title,
                    doc_type=record.doc_type,
                    checksum=record.checksum,
                    content=record.content,
                    content_length=len(record.content),
                    structure_json=None,
                    canonical_id=canonical_id,
                    version=1,
                    is_latest=True,
                )
                session.add(document)
                session.flush()
            else:
                # Same checksum means no content change -- skip re-indexing
                if existing.checksum == record.checksum:
                    return existing.id

                # Content changed: archive the old version and create a new one
                canonical_id = existing.canonical_id or hashlib.sha256(record.source_uri.encode()).hexdigest()[:16]
                old_version = existing.version or 1
                new_version = old_version + 1

                # Mark old document as no longer latest and rename its source_uri to free the unique constraint
                existing.is_latest = False
                existing.canonical_id = canonical_id
                existing.source_uri = f"{record.source_uri}@v{old_version}"
                session.flush()

                # Create new version as a fresh row
                document = Document(
                    source_uri=record.source_uri,
                    title=record.title,
                    doc_type=record.doc_type,
                    checksum=record.checksum,
                    content=record.content,
                    content_length=len(record.content),
                    structure_json=None,
                    canonical_id=canonical_id,
                    version=new_version,
                    is_latest=True,
                )
                session.add(document)
                session.flush()

            for index, chunk in enumerate(record.chunks):
                session.add(
                    DocumentChunk(
                        document_id=document.id,
                        chunk_index=index,
                        content=chunk,
                        token_estimate=max(1, len(chunk) // 4),
                    )
                )
            session.flush()
            return document.id

    def list_documents(self, doc_type: str = "", limit: int = 20, offset: int = 0) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            query = select(Document).order_by(Document.updated_at.desc(), Document.id.desc()).limit(limit).offset(offset)
            if doc_type:
                query = query.where(Document.doc_type == doc_type)
            rows = list(session.scalars(query))
            return [self._document_summary(row) for row in rows]

    def list_documents_for_analysis(self) -> list[dict[str, object]]:
        with session_scope(self.session_factory) as session:
            rows = list(session.scalars(select(Document).order_by(Document.id.asc())))
            return [
                {
                    "id": row.id,
                    "source_uri": row.source_uri,
                    "title": row.title,
                    "doc_type": row.doc_type,
                    "content": row.content,
                }
                for row in rows
            ]

    def store_document_analysis(
        self,
        document_id: int,
        structure: dict[str, object],
        entities: list[tuple[str, int]],
        relationships: list[tuple[str, str, int]],
        claims: list[tuple[str, str, str]] = None,
    ) -> None:
        with session_scope(self.session_factory) as session:
            document = session.scalar(
                select(Document)
                .options(
                    selectinload(Document.entities),
                    selectinload(Document.relationships),
                    selectinload(Document.claims),
                )
                .where(Document.id == document_id)
            )
            if document is None:
                return
            document.structure_json = json.dumps(structure, sort_keys=True)
            for entity in list(document.entities):
                session.delete(entity)
            for relationship in list(document.relationships):
                session.delete(relationship)
            for claim in list(document.claims):
                session.delete(claim)
            session.flush()

            for name, score in entities:
                session.add(DocumentEntity(document_id=document.id, name=name, score=score))
            for source_name, target_name, weight in relationships:
                session.add(
                    DocumentRelationship(
                        document_id=document.id,
                        source_name=source_name,
                        target_name=target_name,
                        weight=weight,
                    )
                )
            if claims:
                for entity_name, claim_text, status in claims:
                    session.add(
                        DocumentClaim(
                            document_id=document.id,
                            entity_name=entity_name,
                            claim=claim_text,
                            status=status,
                        )
                    )

    def replace_communities(self, communities: list[CommunityResult]) -> None:
        with session_scope(self.session_factory) as session:
            for membership in session.scalars(select(DocumentCommunityMembership)):
                session.delete(membership)
            for community in session.scalars(select(DocumentCommunity)):
                session.delete(community)
            session.flush()

            for community in communities:
                row = DocumentCommunity(
                    label=community.label,
                    summary=community.summary,
                    top_terms_json=json.dumps(community.top_terms),
                )
                session.add(row)
                session.flush()
                for document_id in community.document_ids:
                    session.add(DocumentCommunityMembership(community_id=row.id, document_id=document_id))

    def store_chunk_embeddings(self, chunk_embeddings: list[tuple[int, bytes]]) -> None:
        """Store embeddings for chunks."""
        with session_scope(self.session_factory) as session:
            for chunk_id, embedding_bytes in chunk_embeddings:
                chunk = session.scalar(select(DocumentChunk).where(DocumentChunk.id == chunk_id))
                if chunk:
                    chunk.embedding = embedding_bytes

    def store_entity_embeddings(self, entity_embeddings: list[tuple[int, bytes]]) -> None:
        """Store embeddings for entities."""
        with session_scope(self.session_factory) as session:
            for entity_id, embedding_bytes in entity_embeddings:
                entity = session.scalar(select(DocumentEntity).where(DocumentEntity.id == entity_id))
                if entity:
                    entity.embedding = embedding_bytes

    def search_chunks_by_embedding(self, query_embedding: list[float], limit: int = 10, min_score: float = 0.3) -> list[dict]:
        """Search chunks by vector similarity."""
        import struct
        import math

        query_bytes = struct.pack(f"<{len(query_embedding)}f", *query_embedding)
        query_vec = query_embedding

        with session_scope(self.session_factory) as session:
            chunks = list(session.scalars(
                select(DocumentChunk)
                .join(Document)
                .where(DocumentChunk.embedding.isnot(None))
                .order_by(DocumentChunk.id.asc())
            ))

            scored_chunks = []
            for chunk in chunks:
                if not chunk.embedding:
                    continue
                try:
                    chunk_vec = struct.unpack(f"<{len(chunk.embedding)//4}f", chunk.embedding)
                    score = self._cosine_similarity(query_vec, chunk_vec)
                    if score >= min_score:
                        scored_chunks.append((chunk, score))
                except Exception:
                    continue

            scored_chunks.sort(key=lambda x: (-x[1], x[0].id))
            # Build result dicts with document-level version info
            results = []
            doc_cache: dict[int, Document] = {}
            for chunk, score in scored_chunks[:limit]:
                if chunk.document_id not in doc_cache:
                    doc_cache[chunk.document_id] = session.get(Document, chunk.document_id)
                doc = doc_cache[chunk.document_id]
                is_latest = doc.is_latest if doc and doc.is_latest is not None else True
                version_weight = 1.0 if is_latest else 0.7
                results.append({
                    "id": chunk.id,
                    "document_id": chunk.document_id,
                    "content": chunk.content,
                    "score": round(score * version_weight, 4),
                    "is_latest": is_latest,
                })
            return results

    def list_entities(self, document_id: int) -> list[dict]:
        """List entities for a document."""
        with session_scope(self.session_factory) as session:
            entities = list(session.scalars(
                select(DocumentEntity).where(DocumentEntity.document_id == document_id)
            ))
            return [
                {
                    "id": e.id,
                    "name": e.name,
                    "type": e.entity_type,
                    "description": e.description,
                    "score": e.score,
                }
                for e in entities
            ]

    def get_entity(self, entity_id: int) -> dict | None:
        """Get a single entity by ID."""
        with session_scope(self.session_factory) as session:
            entity = session.scalar(select(DocumentEntity).where(DocumentEntity.id == entity_id))
            if not entity:
                return None
            return {
                "id": entity.id,
                "name": entity.name,
                "type": entity.entity_type,
                "description": entity.description,
                "score": entity.score,
            }

    def get_entity_by_name(self, name: str) -> dict | None:
        """Get a single entity by name."""
        with session_scope(self.session_factory) as session:
            entity = session.scalar(select(DocumentEntity).where(DocumentEntity.name == name))
            if not entity:
                return None
            return {
                "id": entity.id,
                "name": entity.name,
                "type": entity.entity_type,
                "description": entity.description,
                "score": entity.score,
            }

    def list_relationships_from_entity(self, source_name: str) -> list[dict]:
        """Get relationships where entity is source."""
        with session_scope(self.session_factory) as session:
            rels = list(session.scalars(
                select(DocumentRelationship).where(DocumentRelationship.source_name == source_name)
            ))
            return [
                {
                    "source": r.source_name,
                    "target": r.target_name,
                    "predicate": r.predicate,
                    "weight": r.weight,
                }
                for r in rels
            ]

    def list_relationships_to_entity(self, target_name: str) -> list[dict]:
        """Get relationships where entity is target."""
        with session_scope(self.session_factory) as session:
            rels = list(session.scalars(
                select(DocumentRelationship).where(DocumentRelationship.target_name == target_name)
            ))
            return [
                {
                    "source": r.source_name,
                    "target": r.target_name,
                    "predicate": r.predicate,
                    "weight": r.weight,
                }
                for r in rels
            ]

    def list_entities_in_community(self, community_id: int) -> list[dict]:
        """List entities in a community."""
        with session_scope(self.session_factory) as session:
            community = session.scalar(
                select(DocumentCommunity)
                .options(selectinload(DocumentCommunity.members))
                .where(DocumentCommunity.id == community_id)
            )
            if not community:
                return []

            doc_ids = [m.document_id for m in community.members]
            entities = list(session.scalars(
                select(DocumentEntity).where(DocumentEntity.document_id.in_(doc_ids))
            ))
            return [
                {
                    "id": e.id,
                    "name": e.name,
                    "type": e.entity_type,
                    "description": e.description,
                    "score": e.score,
                }
                for e in entities
            ]

    @staticmethod
    def _cosine_similarity(a: list[float], b: list[float]) -> float:
        """Compute cosine similarity."""
        import math
        if len(a) != len(b) or not a:
            return 0.0
        dot = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(x * x for x in b))
        if norm_a == 0 or norm_b == 0:
            return 0.0
        return dot / (norm_a * norm_b)

    def get_document(self, document_id: int) -> dict[str, object] | None:
        with session_scope(self.session_factory) as session:
            row = session.scalar(
                select(Document)
                .options(
                    selectinload(Document.chunks),
                    selectinload(Document.entities),
                    selectinload(Document.relationships),
                )
                .where(Document.id == document_id)
            )
            if row is None:
                return None
            payload = self._document_summary(row)
            payload["content"] = row.content
            payload["structure"] = json.loads(row.structure_json) if row.structure_json else None
            payload["entities"] = [{"name": entity.name, "score": entity.score} for entity in row.entities]
            payload["relationships"] = [
                {
                    "source_name": relationship.source_name,
                    "target_name": relationship.target_name,
                    "weight": relationship.weight,
                }
                for relationship in row.relationships
            ]
            payload["chunks"] = [
                {"id": chunk.id, "index": chunk.chunk_index, "content": chunk.content}
                for chunk in sorted(row.chunks, key=lambda item: item.chunk_index)
            ]
            return payload

    def search(self, query_text: str, top_k: int = 5, doc_type: str = "") -> list[SearchResult]:
        terms = [term.lower() for term in query_text.split() if term.strip()]
        if not terms:
            return []

        with session_scope(self.session_factory) as session:
            query = select(DocumentChunk, Document).join(Document, DocumentChunk.document_id == Document.id)
            if doc_type:
                query = query.where(Document.doc_type == doc_type)

            results: list[SearchResult] = []
            for chunk, document in session.execute(query):
                lowered = chunk.content.lower()
                hits = sum(lowered.count(term) for term in terms)
                if hits <= 0:
                    continue
                snippet = chunk.content.strip().replace("\n", " ")
                if len(snippet) > 240:
                    snippet = snippet[:237].rstrip() + "..."
                relevance_score = hits + (1.0 / (1 + chunk.chunk_index))
                version_weight = 1.0 if document.is_latest else 0.7
                score = relevance_score * version_weight
                results.append(
                    SearchResult(
                        document_id=document.id,
                        source_uri=document.source_uri,
                        title=document.title,
                        doc_type=document.doc_type,
                        chunk_id=chunk.id,
                        chunk_index=chunk.chunk_index,
                        snippet=snippet,
                        score=round(score, 4),
                        is_latest=document.is_latest if document.is_latest is not None else True,
                    )
                )

        results.sort(key=lambda item: (-item.score, item.title, item.chunk_index))
        return results[:top_k]

    def list_communities(self) -> list[CommunitySummary]:
        with session_scope(self.session_factory) as session:
            rows = list(
                session.scalars(
                    select(DocumentCommunity)
                    .options(selectinload(DocumentCommunity.members))
                    .order_by(DocumentCommunity.id.asc())
                )
            )
            return [self._community_summary(row) for row in rows]

    def get_community(self, community_id: int) -> CommunitySummary | None:
        with session_scope(self.session_factory) as session:
            row = session.scalar(
                select(DocumentCommunity)
                .options(selectinload(DocumentCommunity.members))
                .where(DocumentCommunity.id == community_id)
            )
            if row is None:
                return None
            return self._community_summary(row)

    def get_stats(self) -> DocumentStats:
        with session_scope(self.session_factory) as session:
            document_count = session.scalar(select(func.count()).select_from(Document)) or 0
            chunk_count = session.scalar(select(func.count()).select_from(DocumentChunk)) or 0
            entity_count = session.scalar(select(func.count()).select_from(DocumentEntity)) or 0
            relationship_count = session.scalar(select(func.count()).select_from(DocumentRelationship)) or 0
            claim_count = session.scalar(select(func.count()).select_from(DocumentClaim)) or 0
            community_count = session.scalar(select(func.count()).select_from(DocumentCommunity)) or 0
            return DocumentStats(
                documents=int(document_count),
                chunks=int(chunk_count),
                entities=int(entity_count),
                relationships=int(relationship_count),
                claims=int(claim_count),
                communities=int(community_count),
            )

    @staticmethod
    def _document_summary(row: Document) -> dict[str, object]:
        provider_status = None
        provider_name = None
        if row.structure_json:
            try:
                structure = json.loads(row.structure_json)
            except json.JSONDecodeError:
                structure = {}
            provider_meta = structure.get("provider") if isinstance(structure, dict) else None
            if isinstance(provider_meta, dict):
                provider_status = provider_meta.get("status")
                provider_name = provider_meta.get("used_provider")
        return {
            "id": row.id,
            "source_uri": row.source_uri,
            "title": row.title,
            "doc_type": row.doc_type,
            "content_length": row.content_length,
            "canonical_id": row.canonical_id,
            "version": row.version or 1,
            "is_latest": row.is_latest if row.is_latest is not None else True,
            "provider_status": provider_status,
            "provider_name": provider_name,
            "updated_at": row.updated_at.isoformat() if row.updated_at else None,
        }

    @staticmethod
    def _community_summary(row: DocumentCommunity) -> CommunitySummary:
        document_ids = sorted(member.document_id for member in row.members)
        return CommunitySummary(
            id=row.id,
            label=row.label,
            summary=row.summary,
            top_terms=json.loads(row.top_terms_json or "[]"),
            document_ids=document_ids,
        )

    def list_all_entities(self, limit: int = 50, offset: int = 0) -> list[dict]:
        """List all entities across all documents, ordered by score desc."""
        with session_scope(self.session_factory) as session:
            rows = list(session.scalars(
                select(DocumentEntity)
                .order_by(DocumentEntity.score.desc(), DocumentEntity.id.asc())
                .limit(limit)
                .offset(offset)
            ))
            return [
                {
                    "id": e.id,
                    "document_id": e.document_id,
                    "name": e.name,
                    "type": e.entity_type,
                    "description": e.description,
                    "score": e.score,
                }
                for e in rows
            ]

    def get_graph_neighborhood(
        self, entity_name: str, depth: int = 2, max_nodes: int = 50
    ) -> dict | None:
        """BFS walk from entity_name, returning nodes and edges up to depth hops."""
        with session_scope(self.session_factory) as session:
            seed = session.scalar(
                select(DocumentEntity).where(DocumentEntity.name == entity_name)
            )
            if seed is None:
                return None

            visited_names: set[str] = set()
            nodes: list[dict] = []
            edges: list[dict] = []
            queue: list[tuple[str, int]] = [(entity_name, 0)]

            while queue and len(nodes) < max_nodes:
                current_name, current_depth = queue.pop(0)
                if current_name in visited_names:
                    continue
                visited_names.add(current_name)

                entity = session.scalar(
                    select(DocumentEntity).where(DocumentEntity.name == current_name)
                )
                nodes.append({
                    "name": current_name,
                    "type": entity.entity_type if entity else "unknown",
                    "description": entity.description if entity else "",
                    "score": entity.score if entity else 0,
                })

                if current_depth >= depth:
                    continue

                # outgoing edges
                rels = list(session.scalars(
                    select(DocumentRelationship)
                    .where(DocumentRelationship.source_name == current_name)
                ))
                for rel in rels:
                    edges.append({
                        "source": rel.source_name,
                        "target": rel.target_name,
                        "predicate": rel.predicate,
                        "weight": rel.weight,
                    })
                    if rel.target_name not in visited_names:
                        queue.append((rel.target_name, current_depth + 1))

                # incoming edges
                in_rels = list(session.scalars(
                    select(DocumentRelationship)
                    .where(DocumentRelationship.target_name == current_name)
                ))
                for rel in in_rels:
                    edges.append({
                        "source": rel.source_name,
                        "target": rel.target_name,
                        "predicate": rel.predicate,
                        "weight": rel.weight,
                    })
                    if rel.source_name not in visited_names:
                        queue.append((rel.source_name, current_depth + 1))

            return {
                "center": entity_name,
                "depth": depth,
                "nodes": nodes[:max_nodes],
                "edges": edges,
            }

    def list_document_versions(self, document_id: int) -> list[dict[str, object]]:
        """Return version information for a document (grouped by canonical_id)."""
        with session_scope(self.session_factory) as session:
            doc = session.get(Document, document_id)
            if doc is None:
                return []
            if doc.canonical_id:
                rows = list(session.scalars(
                    select(Document)
                    .where(Document.canonical_id == doc.canonical_id)
                    .order_by(Document.version.desc())
                ))
            else:
                rows = [doc]
            return [
                {
                    "id": row.id,
                    "source_uri": row.source_uri,
                    "version": row.version or 1,
                    "is_latest": row.is_latest if row.is_latest is not None else True,
                    "checksum": row.checksum,
                    "canonical_id": row.canonical_id,
                    "updated_at": row.updated_at.isoformat() if row.updated_at else None,
                }
                for row in rows
            ]
