"""Post-ingest observation categorizer using Nemotron Nano.

Classifies observations into domain tags and validates category assignments.
Designed as a post-ingest pipeline stage: observations are created first with
whatever metadata the caller provides, then this module enriches them with
LLM-derived domain tags.

Usage:
    from greymatter_core.categorizer import ObservationCategorizer
    from greymatter_core.llm_gateway import OllamaLLM

    categorizer = ObservationCategorizer(OllamaLLM())
    result = categorizer.categorize(observation_dict)
    # result: {"domains": ["python", "testing"], "suggested_category": "pattern", ...}

    # Batch mode: backfill untagged observations
    stats = categorizer.categorize_batch(db, batch_size=50)
"""
from __future__ import annotations

import logging
from typing import Any, TYPE_CHECKING

if TYPE_CHECKING:
    from greymatter_core.db import Database
    from greymatter_core.llm_gateway import OllamaLLM

logger = logging.getLogger(__name__)

VALID_CATEGORIES = frozenset({
    "decision", "correction", "pattern", "insight", "reference", "gotcha",
})

DOMAIN_TAXONOMY = [
    "architecture", "database", "security", "testing", "deployment",
    "performance", "api", "frontend", "backend", "infrastructure",
    "devops", "monitoring", "documentation", "python", "rust", "go",
    "javascript", "typescript", "swift", "networking", "authentication",
    "authorization", "caching", "search", "ml", "ai", "data-pipeline",
    "configuration", "error-handling", "concurrency", "business",
    "pricing", "ux", "accessibility", "compliance", "observability",
]

_CATEGORIZE_SYSTEM = (
    "You are a domain categorizer for a developer knowledge base. "
    "Given an observation (title + content + category), assign 1-3 domain tags.\n\n"
    "Domain tags should be lowercase, hyphenated if multi-word, and concise.\n"
    "Prefer tags from this taxonomy when they fit:\n"
    + ", ".join(DOMAIN_TAXONOMY) + "\n\n"
    "You may use tags outside the taxonomy if none fit.\n\n"
    "Also validate the category. If the assigned category seems wrong, "
    "suggest a better one.\n\n"
    "Return JSON:\n"
    "{\n"
    '  "domains": ["tag1", "tag2"],\n'
    '  "suggested_category": "pattern",\n'
    '  "category_changed": false,\n'
    '  "confidence": 0.85\n'
    "}\n\n"
    "Rules:\n"
    "- domains: 1-3 tags, lowercase, no spaces\n"
    "- suggested_category: one of: decision, correction, pattern, insight, reference, gotcha\n"
    "- category_changed: true only if you think the original category is wrong\n"
    "- confidence: float 0.0-1.0, how confident you are in the domain assignment"
)

_BATCH_CATEGORIZE_SYSTEM = (
    "You are a domain categorizer for a developer knowledge base. "
    "Given a batch of observations, assign domain tags to each one.\n\n"
    "Domain tags should be lowercase, hyphenated if multi-word, and concise.\n"
    "Prefer tags from this taxonomy when they fit:\n"
    + ", ".join(DOMAIN_TAXONOMY) + "\n\n"
    "Return a JSON array where each element has:\n"
    "{\n"
    '  "id": "the observation id",\n'
    '  "domains": ["tag1", "tag2"],\n'
    '  "suggested_category": "pattern",\n'
    '  "category_changed": false,\n'
    '  "confidence": 0.85\n'
    "}\n\n"
    "Rules:\n"
    "- domains: 1-3 tags per observation, lowercase, no spaces\n"
    "- suggested_category: one of: decision, correction, pattern, insight, reference, gotcha\n"
    "- category_changed: true only if the original category seems wrong\n"
    "- confidence: float 0.0-1.0 per observation\n"
    "- Return results in the same order as the input"
)


class ObservationCategorizer:
    """Enriches observations with LLM-derived domain tags.

    Uses Nemotron Nano via OllamaLLM to classify observations into domain
    categories. Operates as a post-ingest pipeline stage.

    Args:
        llm: OllamaLLM instance (defaults to nemotron-3-nano).
        min_confidence: Minimum LLM confidence to accept domain tags.
        auto_recategorize: Whether to accept category changes from the LLM.
    """

    def __init__(
        self,
        llm: OllamaLLM,
        min_confidence: float = 0.5,
        auto_recategorize: bool = False,
    ):
        self.llm = llm
        self.min_confidence = min_confidence
        self.auto_recategorize = auto_recategorize

    def categorize(self, observation: dict[str, Any]) -> dict[str, Any]:
        """Categorize a single observation."""
        title = observation.get("title", "")
        content = observation.get("content", "")
        category = observation.get("category", "insight")

        if not content and not title:
            return _empty_result(category)

        prompt = (
            f"Observation:\n"
            f"  Title: {title}\n"
            f"  Content: {content}\n"
            f"  Category: {category}\n\n"
            f"Assign domain tags and validate the category."
        )

        try:
            result = self.llm.generate_json(prompt, system=_CATEGORIZE_SYSTEM)
        except Exception:
            logger.debug("LLM categorization failed, returning empty result")
            return _empty_result(category)

        return _parse_result(result, category)

    def categorize_and_apply(
        self,
        observation: dict[str, Any],
        db: Database,
    ) -> dict[str, Any]:
        """Categorize an observation and write the results to the database."""
        result = self.categorize(observation)
        obs_id = observation.get("id", "")

        if not obs_id or not result["domains"]:
            return result

        domains_str = ",".join(result["domains"])

        if result["category_changed"] and self.auto_recategorize:
            db.execute(
                "UPDATE observations SET domains = ?, category = ?, "
                "updated_at = datetime('now') WHERE id = ?",
                (domains_str, result["suggested_category"], obs_id),
            )
            logger.info(
                "Categorized %s: domains=%s, category=%s->%s",
                obs_id, domains_str, observation["category"],
                result["suggested_category"],
            )
        else:
            db.execute(
                "UPDATE observations SET domains = ?, "
                "updated_at = datetime('now') WHERE id = ?",
                (domains_str, obs_id),
            )
            logger.info("Categorized %s: domains=%s", obs_id, domains_str)

        return result

    def categorize_batch(
        self,
        db: Database,
        batch_size: int = 50,
        untagged_only: bool = True,
    ) -> dict[str, Any]:
        """Batch-categorize observations in the database."""
        where = "pruned = 0"
        if untagged_only:
            where += " AND (domains IS NULL OR domains = '')"

        rows = db.execute(
            f"SELECT id, title, content, category, domains FROM observations "
            f"WHERE {where} ORDER BY created_at DESC LIMIT ?",
            (batch_size,),
        )

        stats = {"processed": 0, "tagged": 0, "failed": 0, "skipped": 0}

        if not rows:
            return stats

        for i in range(0, len(rows), 10):
            chunk = rows[i:i + 10]
            chunk_results = self._categorize_chunk(chunk)

            for obs, result in zip(chunk, chunk_results):
                stats["processed"] += 1

                if not result["domains"]:
                    stats["skipped"] += 1
                    continue

                if result["confidence"] < self.min_confidence:
                    stats["skipped"] += 1
                    continue

                domains_str = ",".join(result["domains"])

                try:
                    if result["category_changed"] and self.auto_recategorize:
                        db.execute(
                            "UPDATE observations SET domains = ?, category = ?, "
                            "updated_at = datetime('now') WHERE id = ?",
                            (domains_str, result["suggested_category"], obs["id"]),
                        )
                    else:
                        db.execute(
                            "UPDATE observations SET domains = ?, "
                            "updated_at = datetime('now') WHERE id = ?",
                            (domains_str, obs["id"]),
                        )
                    stats["tagged"] += 1
                except Exception:
                    logger.exception("Failed to update observation %s", obs["id"])
                    stats["failed"] += 1

        logger.info(
            "Batch categorization complete: %d processed, %d tagged, "
            "%d failed, %d skipped",
            stats["processed"], stats["tagged"],
            stats["failed"], stats["skipped"],
        )
        return stats

    def _categorize_chunk(
        self, observations: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Categorize a chunk of observations in a single LLM call."""
        if len(observations) == 1:
            return [self.categorize(observations[0])]

        lines = []
        for obs in observations:
            lines.append(
                f"- ID: {obs['id']}\n"
                f"  Title: {obs.get('title', '')}\n"
                f"  Content: {obs.get('content', '')}\n"
                f"  Category: {obs.get('category', 'insight')}"
            )
        prompt = "Categorize these observations:\n\n" + "\n\n".join(lines)

        try:
            results = self.llm.generate_json(
                prompt, system=_BATCH_CATEGORIZE_SYSTEM
            )
            if not isinstance(results, list) or len(results) != len(observations):
                raise ValueError("Batch result count mismatch")

            parsed = []
            for r, obs in zip(results, observations):
                parsed.append(_parse_result(r, obs.get("category", "insight")))
            return parsed

        except Exception:
            logger.debug("Batch categorization failed, falling back to individual")
            return [self.categorize(obs) for obs in observations]


def _parse_result(
    result: Any, original_category: str
) -> dict[str, Any]:
    """Parse and validate a categorization result from the LLM."""
    if not isinstance(result, dict):
        return _empty_result(original_category)

    raw_domains = result.get("domains", [])
    if not isinstance(raw_domains, list):
        raw_domains = []

    domains = []
    for d in raw_domains:
        tag = str(d).strip().lower().replace(" ", "-")
        if tag and len(tag) <= 40:
            domains.append(tag)
    domains = domains[:3]

    suggested = result.get("suggested_category", original_category)
    if suggested not in VALID_CATEGORIES:
        suggested = original_category

    category_changed = (
        result.get("category_changed", False) is True
        and suggested != original_category
    )

    raw_conf = result.get("confidence", 0.5)
    try:
        confidence = min(max(float(raw_conf), 0.0), 1.0)
    except (TypeError, ValueError):
        confidence = 0.5

    return {
        "domains": domains,
        "suggested_category": suggested,
        "category_changed": category_changed,
        "confidence": confidence,
    }


def _empty_result(category: str) -> dict[str, Any]:
    """Return an empty categorization result."""
    return {
        "domains": [],
        "suggested_category": category,
        "category_changed": False,
        "confidence": 0.0,
    }
