"""Local AI-relevance classifier using Ollama.

Two-stage filter to avoid spending Claude tokens on irrelevant drafts:
1. Embedding similarity — fast cosine check against a reference description
2. Chat classification — small local model for borderline cases

Both stages run locally via Ollama (zero cost).
"""

from __future__ import annotations

import numpy as np
import ollama as ollama_lib
from rich.console import Console

from .config import Config

console = Console()

# Reference description of what we're looking for.
# Embedding of this text is compared against each draft's abstract.
REFERENCE_DESCRIPTION = """
AI agent protocols, autonomous agent communication, agent-to-agent interaction,
agent identity and authentication, agent authorization, agent discovery,
large language model integration with network protocols, agentic systems,
machine learning for network operations, AI safety in networked systems,
model context protocol, multi-agent coordination, agent task delegation,
generative AI infrastructure, intelligent network automation,
trustworthy AI systems, AI governance in standards.
"""

# Thresholds for the two-stage filter (calibrated against 434 drafts + 73 FPs)
# TP avg similarity: 0.685, FP avg: 0.598
SIMILARITY_ACCEPT = 0.72   # Above this: definitely relevant, skip chat
SIMILARITY_REJECT = 0.50   # Below this: definitely irrelevant, skip chat
# Between REJECT and ACCEPT: borderline, use chat model to decide

CLASSIFY_PROMPT = """\
You are classifying IETF Internet-Drafts for an AI/agent standards tracker.

A draft is RELEVANT if it relates to ANY of these topics:
- AI agents, autonomous agents, multi-agent systems
- Agent identity, authentication, authorization, discovery
- Agent-to-agent (A2A) communication protocols
- Large language models (LLMs), generative AI
- Machine learning in network operations
- AI safety, alignment, trustworthiness
- Model Context Protocol (MCP), agentic workflows
- OAuth/JWT/credentials for agents or AI systems
- Autonomous network operations using AI
- Intelligent network management or traffic handling

A draft is NOT relevant if it only covers:
- Pure cryptography without AI/agent context
- General networking protocols (BGP, DNS, TLS) without AI
- Email, HTTP, or web standards without AI/agent features

Title: {title}

Abstract: {abstract}

Is this draft relevant to AI agents or related topics? Answer ONLY "yes" or "no"."""


class Classifier:
    """Classify drafts as AI-relevant using local Ollama models."""

    def __init__(self, config: Config | None = None):
        self.config = config or Config.load()
        self.client = ollama_lib.Client(host=self.config.ollama_url)
        self._ref_embedding: np.ndarray | None = None

    def close(self) -> None:
        if hasattr(self.client, '_client'):
            self.client._client.close()

    def __enter__(self):
        return self

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

    def _get_reference_embedding(self) -> np.ndarray:
        """Get (cached) embedding of the reference AI description."""
        if self._ref_embedding is None:
            resp = self.client.embed(
                model=self.config.ollama_embed_model,
                input=REFERENCE_DESCRIPTION.strip(),
            )
            self._ref_embedding = np.array(resp["embeddings"][0], dtype=np.float32)
        return self._ref_embedding

    def _embed(self, text: str) -> np.ndarray:
        """Embed a text string."""
        resp = self.client.embed(
            model=self.config.ollama_embed_model,
            input=text[:8000],
        )
        return np.array(resp["embeddings"][0], dtype=np.float32)

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

    def _chat_classify(self, title: str, abstract: str) -> bool:
        """Ask local chat model whether a draft is AI-related."""
        prompt = CLASSIFY_PROMPT.format(title=title, abstract=abstract[:2000])
        try:
            resp = self.client.chat(
                model=self.config.ollama_classify_model,
                messages=[{"role": "user", "content": prompt}],
                options={"temperature": 0.0, "num_predict": 10},
            )
            answer = resp["message"]["content"].strip().lower()
            return answer.startswith("yes")
        except Exception as e:
            console.print(f"[yellow]Chat classify failed: {e}, defaulting to relevant[/yellow]")
            return True  # err on the side of inclusion

    def classify(self, title: str, abstract: str) -> tuple[bool, float, str]:
        """Classify a draft as AI-relevant.

        Returns:
            (is_relevant, similarity_score, method)
            method is one of: "embedding_accept", "embedding_reject", "chat_yes", "chat_no"
        """
        text = f"{title}\n{abstract}"
        ref = self._get_reference_embedding()
        emb = self._embed(text)
        sim = self._cosine_similarity(emb, ref)

        if sim >= SIMILARITY_ACCEPT:
            return True, sim, "embedding_accept"
        if sim <= SIMILARITY_REJECT:
            return False, sim, "embedding_reject"

        # Borderline — ask chat model
        is_relevant = self._chat_classify(title, abstract)
        method = "chat_yes" if is_relevant else "chat_no"
        return is_relevant, sim, method

    def classify_batch(
        self, drafts: list[dict], verbose: bool = True
    ) -> tuple[list[dict], list[dict]]:
        """Classify a batch of drafts.

        Args:
            drafts: list of dicts with at least 'name', 'title', 'abstract' keys

        Returns:
            (relevant, irrelevant) — two lists of draft dicts
        """
        relevant = []
        irrelevant = []
        stats = {"embedding_accept": 0, "embedding_reject": 0, "chat_yes": 0, "chat_no": 0}

        for i, d in enumerate(drafts):
            is_rel, sim, method = self.classify(
                d.get("title", ""), d.get("abstract", "")
            )
            stats[method] += 1

            if verbose and (i + 1) % 10 == 0:
                console.print(f"  Classified {i + 1}/{len(drafts)}...")

            if is_rel:
                relevant.append(d)
            else:
                irrelevant.append(d)

        if verbose:
            console.print(
                f"\n  [green]Relevant: {len(relevant)}[/green]  "
                f"[red]Irrelevant: {len(irrelevant)}[/red]\n"
                f"  Embedding accept: {stats['embedding_accept']}  "
                f"  Embedding reject: {stats['embedding_reject']}\n"
                f"  Chat yes: {stats['chat_yes']}  "
                f"  Chat no: {stats['chat_no']}"
            )

        return relevant, irrelevant
