"""
reflection_sidecar.py — Reflection Sidecar (Double-Token Innovation)

Analizza i log di errore di ogni tool call durante la sessione e aggiorna
session_rules.md in tempo reale. Questo file viene iniettato nel system prompt
dell'agente principale per correggere il comportamento on-the-fly.

Architettura:
  Token A (agente principale) → esegue tool, chiama log_error()
  Token B (sidecar critic)    → analizza pattern, scrive regole
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import re
import time
from collections import defaultdict, deque
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

_logger = logging.getLogger("agente_ai.reflection_sidecar")

# ── Config ────────────────────────────────────────────────────────────────────

_RULES_FILE = Path(os.getenv("SIDECAR_RULES_FILE", "/data/session_rules.md"))
_MAX_ERRORS_BEFORE_REFLECT = int(os.getenv("SIDECAR_REFLECT_THRESHOLD", "2"))
_RULE_TTL_S = int(os.getenv("SIDECAR_RULE_TTL_S", "3600"))  # 1h

# Token B per NVIDIA NIM (Reflection Critic — modello leggero, bassa latenza)
_NVIDIA_API   = "https://integrate.api.nvidia.com/v1"
_CRITIC_MODEL = os.getenv("NVIDIA_B_MODEL", "meta/llama-3.3-70b-instruct")
_NVIDIA_KEY_B = os.getenv("NVIDIA_API_KEY_B", "")


# ── Data model ────────────────────────────────────────────────────────────────

@dataclass
class ErrorEvent:
    tool: str
    error: str
    context: str
    ts: float = field(default_factory=time.monotonic)


@dataclass
class SessionRule:
    pattern: str          # cosa ha causato l'errore (regex / descrizione)
    rule: str             # istruzione correttiva per l'agente
    tool: str
    created_at: float = field(default_factory=time.time)
    hit_count: int = 0


# ── Sidecar core ─────────────────────────────────────────────────────────────

class ReflectionSidecar:
    """
    Singleton per sessione. Riceve errori, li analizza con Token B (NVIDIA),
    aggiorna session_rules.md che viene iniettato nel prompt principale.
    """

    def __init__(self) -> None:
        self._errors: list[ErrorEvent] = []
        self._rules: list[SessionRule] = []
        self._tool_error_counts: dict[str, int] = defaultdict(int)
        self._lock = asyncio.Lock()
        self._reflect_task: asyncio.Task | None = None

    async def log_error(
        self,
        tool: str,
        error: str,
        context: str = "",
    ) -> None:
        """Registra un errore. Se lo stesso tool fallisce >= threshold, avvia reflection."""
        async with self._lock:
            evt = ErrorEvent(tool=tool, error=error[:500], context=context[:300])
            self._errors.append(evt)
            self._tool_error_counts[tool] += 1
            count = self._tool_error_counts[tool]

        if count >= _MAX_ERRORS_BEFORE_REFLECT:
            # Avvia reflection in background (non blocca l'agente principale)
            if self._reflect_task is None or self._reflect_task.done():
                self._reflect_task = asyncio.create_task(
                    self._reflect_and_update(tool, error, context)
                )

    async def _reflect_and_update(
        self, tool: str, last_error: str, context: str
    ) -> None:
        """Token B: analizza gli errori e genera una regola correttiva."""
        if not _NVIDIA_KEY_B:
            _logger.warning("reflection_sidecar: NVIDIA_API_KEY_B non configurato — skip")
            return

        # Aggrega tutti gli errori del tool
        relevant = [e for e in self._errors if e.tool == tool][-5:]
        error_summary = "\n".join(f"- [{e.tool}] {e.error}" for e in relevant)

        prompt = f"""Sei un critico di qualità per un agente AI. Analizza questi errori ripetuti:

TOOL: {tool}
ERRORI:
{error_summary}

CONTESTO ULTIMO ERRORE: {context}

Scrivi UNA regola correttiva concisa (max 2 righe) che l'agente deve seguire per evitare
di ripetere questo errore. Formato: "REGOLA [{tool}]: <istruzione diretta all'agente>"
Rispondi solo con la regola, nessun altro testo."""

        try:
            import urllib.request
            payload = json.dumps({
                "model": _CRITIC_MODEL,
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": 120,
                "temperature": 0.1,
            }).encode()
            req = urllib.request.Request(
                f"{_NVIDIA_API}/chat/completions",
                data=payload,
                headers={
                    "Authorization": f"Bearer {_NVIDIA_KEY_B}",
                    "Content-Type": "application/json",
                },
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=15) as resp:
                data = json.loads(resp.read())
            rule_text = data["choices"][0]["message"]["content"].strip()

            new_rule = SessionRule(
                pattern=last_error[:100],
                rule=rule_text,
                tool=tool,
            )
            async with self._lock:
                # Dedup: rimuovi regole vecchie per lo stesso tool
                self._rules = [r for r in self._rules if r.tool != tool]
                self._rules.append(new_rule)
                self._tool_error_counts[tool] = 0  # reset counter

            await self._write_rules_file()
            _logger.info(f"reflection_sidecar: nuova regola generata per {tool}")

        except Exception as exc:
            _logger.warning(f"reflection_sidecar: reflection fallita — {exc}")

    async def _write_rules_file(self) -> None:
        """Scrive session_rules.md — viene iniettato nel system prompt principale."""
        now = time.time()
        active = [r for r in self._rules if (now - r.created_at) < _RULE_TTL_S]
        if not active:
            return
        lines = ["# Session Rules (auto-generate dal Reflection Sidecar)\n"]
        lines += [f"- {r.rule}" for r in active]
        lines.append(f"\n_Aggiornato: {time.strftime('%H:%M:%S')}_")
        try:
            _RULES_FILE.parent.mkdir(parents=True, exist_ok=True)
            _RULES_FILE.write_text("\n".join(lines), encoding="utf-8")
        except Exception as exc:
            _logger.warning(f"reflection_sidecar: scrittura rules file fallita — {exc}")

    def get_rules_for_prompt(self) -> str:
        """Legge session_rules.md per l'iniezione nel system prompt."""
        try:
            if _RULES_FILE.exists():
                content = _RULES_FILE.read_text(encoding="utf-8").strip()
                if content and len(content) > 30:
                    return f"\n\n---\n{content}\n---"
        except Exception:
            pass
        return ""

    def reset(self) -> None:
        """Reset a inizio nuova sessione."""
        self._errors.clear()
        self._rules.clear()
        self._tool_error_counts.clear()
        try:
            _RULES_FILE.unlink(missing_ok=True)
        except Exception:
            pass


# ── Singleton ─────────────────────────────────────────────────────────────────
_sidecar: ReflectionSidecar | None = None


def get_sidecar() -> ReflectionSidecar:
    global _sidecar
    if _sidecar is None:
        _sidecar = ReflectionSidecar()
    return _sidecar


async def log_tool_error(tool: str, error: str, context: str = "") -> None:
    """Shortcut globale — chiamare dopo ogni tool call fallita."""
    await get_sidecar().log_error(tool, error, context)


def get_session_rules() -> str:
    """Shortcut globale — iniettare nel system prompt principale."""
    return get_sidecar().get_rules_for_prompt()