"""Embedding generation via Ollama and similarity computation."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import numpy as np
import ollama as ollama_lib
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, MofNCompleteColumn

from .config import Config
from .db import Database

console = Console()


def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    dot = np.dot(a, b)
    norm = np.linalg.norm(a) * np.linalg.norm(b)
    if norm == 0:
        return 0.0
    return float(dot / norm)


class Embedder:
    def __init__(self, config: Config | None = None, db: Database | None = None):
        self.config = config or Config.load()
        self.db = db or Database(self.config)
        self.client = ollama_lib.Client(host=self.config.ollama_url)

    def close(self) -> None:
        """Close the underlying Ollama HTTP client."""
        if hasattr(self.client, '_client'):
            self.client._client.close()

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()

    def embed_text(self, text: str) -> np.ndarray:
        """Generate an embedding for a single text string."""
        # Truncate to ~8k tokens worth of text (roughly 32k chars)
        truncated = text[:32000]
        resp = self.client.embed(model=self.config.ollama_embed_model, input=truncated)
        return np.array(resp["embeddings"][0], dtype=np.float32)

    def embed_draft(self, draft_name: str) -> np.ndarray | None:
        """Generate and store an embedding for a draft using its abstract + title."""
        draft = self.db.get_draft(draft_name)
        if draft is None:
            console.print(f"[red]Draft not found: {draft_name}[/]")
            return None

        # Combine title + abstract + beginning of full text for richer embedding
        parts = [draft.title, draft.abstract]
        if draft.full_text:
            # Include first ~4k chars of body
            parts.append(draft.full_text[:4000])
        text = "\n\n".join(p for p in parts if p)

        vec = self.embed_text(text)
        self.db.store_embedding(draft_name, self.config.ollama_embed_model, vec)
        return vec

    def embed_all_missing(self) -> int:
        """Generate embeddings for all drafts that don't have one yet."""
        missing = self.db.drafts_without_embeddings()
        if not missing:
            console.print("All drafts already have embeddings.")
            return 0

        count = 0
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            MofNCompleteColumn(),
            console=console,
        ) as progress:
            task = progress.add_task("Generating embeddings...", total=len(missing))
            for name in missing:
                try:
                    self.embed_draft(name)
                    count += 1
                except Exception as e:
                    console.print(f"[red]Failed to embed {name}: {e}[/]")
                progress.advance(task)

        console.print(f"Generated [bold green]{count}[/] embeddings")
        return count

    def find_similar(self, draft_name: str, top_n: int = 10) -> list[tuple[str, float]]:
        """Find the most similar drafts to a given draft."""
        target_vec = self.db.get_embedding(draft_name)
        if target_vec is None:
            # Try generating it on the fly
            target_vec = self.embed_draft(draft_name)
            if target_vec is None:
                return []

        all_embeddings = self.db.all_embeddings()
        similarities: list[tuple[str, float]] = []
        for name, vec in all_embeddings.items():
            if name == draft_name:
                continue
            sim = _cosine_similarity(target_vec, vec)
            similarities.append((name, sim))

        similarities.sort(key=lambda x: x[1], reverse=True)
        return similarities[:top_n]

    def similarity_matrix(self) -> tuple[list[str], np.ndarray]:
        """Compute pairwise similarity matrix for all embedded drafts.

        Uses a file-based cache keyed by the hash of embedding draft names.
        If the set of embedded drafts hasn't changed, the cached matrix is
        reloaded from disk instead of recomputing O(n^2) cosine similarities.
        """
        all_embeddings = self.db.all_embeddings()
        names = sorted(all_embeddings.keys())
        n = len(names)

        # Build cache key from sorted draft names
        names_hash = hashlib.sha256("\n".join(names).encode()).hexdigest()[:16]
        cache_dir = Path(self.config.db_path).parent / ".cache"
        cache_meta = cache_dir / f"sim_matrix_{names_hash}.json"
        cache_npy = cache_dir / f"sim_matrix_{names_hash}.npy"

        # Try loading from cache
        if cache_meta.exists() and cache_npy.exists():
            try:
                cached_names = json.loads(cache_meta.read_text())
                if cached_names == names:
                    matrix = np.load(cache_npy)
                    if matrix.shape == (n, n):
                        return names, matrix
            except Exception:
                pass  # Cache corrupted, recompute

        # Compute fresh
        matrix = np.zeros((n, n), dtype=np.float32)
        for i in range(n):
            for j in range(i, n):
                sim = _cosine_similarity(all_embeddings[names[i]], all_embeddings[names[j]])
                matrix[i, j] = sim
                matrix[j, i] = sim

        # Save to cache
        try:
            cache_dir.mkdir(exist_ok=True)
            np.save(cache_npy, matrix)
            cache_meta.write_text(json.dumps(names))
        except Exception:
            pass  # Non-fatal if caching fails

        return names, matrix

    def find_clusters(self, threshold: float = 0.85) -> list[list[str]]:
        """Find clusters of highly similar drafts using simple greedy clustering.

        Methodology notes:
        - Uses greedy single-linkage clustering: once a draft joins a cluster,
          all drafts similar to *it* (but not necessarily to the seed) can join
          too. This can produce "chaining" where semantically distant drafts
          end up in the same cluster through intermediaries.
        - The 0.85 default threshold is an EMPIRICAL CHOICE, not derived from
          a principled analysis. It was selected by manual inspection of draft
          pairs at various thresholds: 0.80 produced too many false positive
          groupings, 0.90 missed obvious topical clusters, and 0.85 yielded
          groups that looked reasonable on spot-checking. A sensitivity analysis
          (running at 0.80, 0.85, 0.90) would strengthen confidence in this
          threshold. The companion threshold of 0.90 used elsewhere for
          "near-duplicates" and 0.98 for "functionally identical" are similarly
          empirical.
        - The embedding model (nomic-embed-text) is a general-purpose model,
          not fine-tuned for technical/standards document similarity. Domain-
          specific embeddings might produce different cluster structures.
        - No comparison to alternative clustering methods (k-means, DBSCAN,
          hierarchical) has been performed.
        """
        names, matrix = self.similarity_matrix()
        if len(names) == 0:
            return []

        visited = set()
        clusters: list[list[str]] = []

        for i, name in enumerate(names):
            if name in visited:
                continue
            cluster = [name]
            visited.add(name)
            for j in range(len(names)):
                if names[j] not in visited and matrix[i, j] >= threshold:
                    cluster.append(names[j])
                    visited.add(names[j])
            if len(cluster) > 1:
                clusters.append(cluster)

        return clusters
