# Copyright 2025 Vijaykumar Singh <singhvjd@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""RL learner for per-provider grounding/hallucination threshold optimization.

This learner uses Thompson Sampling (Bayesian optimization) to learn optimal
confidence thresholds for hallucination detection per provider.

Problem:
- Static thresholds cause false positives (overly strict) or missed
  hallucinations (too lenient)
- Different providers have different hallucination patterns

Strategy:
- Context: (provider, model, response_type)
- Action: Select threshold from discretized range [0.5, 0.95]
- Reward: +0.1 correct, -1.0 false positive, -2.0 false negative

Thompson Sampling is ideal here because:
1. Good exploration/exploitation balance for continuous optimization
2. Naturally handles uncertainty
3. Works well with sparse rewards

Sprint 3: Cache & Grounding Learners
"""

import logging
import math
import random
from datetime import datetime
from typing import Any, Dict, Optional, Tuple

from victor.agent.rl.base import BaseLearner, RLOutcome, RLRecommendation

logger = logging.getLogger(__name__)


class GroundingThresholdLearner(BaseLearner):
    """Learn optimal grounding thresholds using Thompson Sampling.

    Uses Beta distribution conjugate prior for threshold selection.
    Each (provider, response_type) pair maintains its own Beta parameters.

    Attributes:
        name: Always "grounding_threshold"
        db: SQLite database connection
        learning_rate: Update rate for parameters, default 0.1
    """

    # Discretized threshold levels
    THRESHOLD_LEVELS = [0.50, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85, 0.90, 0.95]

    # Default threshold when no data available
    DEFAULT_THRESHOLD = 0.70

    # Response types for state categorization
    RESPONSE_TYPES = ["code_generation", "explanation", "analysis", "edit", "search", "general"]

    # Minimum observations before confident recommendation
    MIN_SAMPLES_FOR_CONFIDENCE = 10

    # Beta prior parameters (weakly informative)
    PRIOR_ALPHA = 1.0
    PRIOR_BETA = 1.0

    def __init__(
        self,
        name: str,
        db_connection: Any,
        learning_rate: float = 0.1,
        provider_adapter: Optional[Any] = None,
    ):
        """Initialize grounding threshold learner.

        Args:
            name: Learner name (should be "grounding_threshold")
            db_connection: SQLite database connection
            learning_rate: Parameter update rate (default 0.1)
            provider_adapter: Optional provider adapter
        """
        super().__init__(name, db_connection, learning_rate, provider_adapter)

        # Beta distribution parameters: (alpha, beta) for each (provider, response_type, threshold)
        # Alpha = successes + prior, Beta = failures + prior
        self._beta_params: Dict[str, Dict[float, Tuple[float, float]]] = {}

        # Track outcomes for analysis
        self._fp_rates: Dict[str, float] = {}  # False positive rates per provider
        self._fn_rates: Dict[str, float] = {}  # False negative rates per provider
        self._total_decisions: int = 0

        # Load state from database
        self._load_state()

    def _ensure_tables(self) -> None:
        """Create tables for grounding threshold learning."""
        cursor = self.db.cursor()

        # Beta parameters table
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS grounding_threshold_params (
                context_key TEXT NOT NULL,
                threshold REAL NOT NULL,
                alpha REAL NOT NULL DEFAULT 1.0,
                beta REAL NOT NULL DEFAULT 1.0,
                sample_count INTEGER NOT NULL DEFAULT 0,
                last_updated TEXT NOT NULL,
                PRIMARY KEY (context_key, threshold)
            )
            """
        )

        # Provider statistics
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS grounding_threshold_stats (
                provider TEXT PRIMARY KEY,
                true_positives INTEGER NOT NULL DEFAULT 0,
                true_negatives INTEGER NOT NULL DEFAULT 0,
                false_positives INTEGER NOT NULL DEFAULT 0,
                false_negatives INTEGER NOT NULL DEFAULT 0,
                last_updated TEXT NOT NULL
            )
            """
        )

        # Decision history
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS grounding_threshold_history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                context_key TEXT NOT NULL,
                threshold REAL NOT NULL,
                result_type TEXT NOT NULL,
                reward REAL NOT NULL,
                timestamp TEXT NOT NULL
            )
            """
        )

        # Indexes
        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_grounding_threshold_context
            ON grounding_threshold_params(context_key)
            """
        )

        self.db.commit()
        logger.debug("RL: grounding_threshold tables ensured")

    def _load_state(self) -> None:
        """Load state from database."""
        cursor = self.db.cursor()

        # Load Beta parameters
        try:
            cursor.execute("SELECT * FROM grounding_threshold_params")
            for row in cursor.fetchall():
                row_dict = dict(row)
                context_key = row_dict["context_key"]
                threshold = row_dict["threshold"]

                if context_key not in self._beta_params:
                    self._beta_params[context_key] = {}

                self._beta_params[context_key][threshold] = (
                    row_dict["alpha"],
                    row_dict["beta"],
                )
                self._total_decisions += row_dict["sample_count"]

        except Exception as e:
            logger.debug(f"RL: Could not load Beta parameters: {e}")

        # Load error rates
        try:
            cursor.execute("SELECT * FROM grounding_threshold_stats")
            for row in cursor.fetchall():
                row_dict = dict(row)
                provider = row_dict["provider"]
                total = (
                    row_dict["true_positives"]
                    + row_dict["true_negatives"]
                    + row_dict["false_positives"]
                    + row_dict["false_negatives"]
                )
                if total > 0:
                    self._fp_rates[provider] = row_dict["false_positives"] / total
                    self._fn_rates[provider] = row_dict["false_negatives"] / total

        except Exception as e:
            logger.debug(f"RL: Could not load error rates: {e}")

        if self._beta_params:
            logger.info(
                f"RL: Loaded {len(self._beta_params)} grounding threshold contexts from database"
            )

    def record_outcome(self, outcome: RLOutcome) -> None:
        """Record grounding verification outcome.

        Expected metadata:
        - provider: Provider name
        - response_type: Type of response (code_generation, explanation, etc.)
        - threshold_used: The threshold that was used
        - result_type: "tp" (true positive), "tn" (true negative),
                       "fp" (false positive), "fn" (false negative)
        - actual_hallucination: Whether there was actually a hallucination
        - detected_hallucination: Whether hallucination was detected

        Args:
            outcome: Outcome with grounding verification data
        """
        provider = outcome.provider
        response_type = outcome.metadata.get("response_type", "general")
        threshold_used = outcome.metadata.get("threshold_used", self.DEFAULT_THRESHOLD)
        result_type = outcome.metadata.get("result_type")

        if not result_type:
            # Infer from actual vs detected
            actual = outcome.metadata.get("actual_hallucination", False)
            detected = outcome.metadata.get("detected_hallucination", False)
            if actual and detected:
                result_type = "tp"  # True positive - correctly detected
            elif not actual and not detected:
                result_type = "tn"  # True negative - correctly passed
            elif not actual and detected:
                result_type = "fp"  # False positive - wrongly flagged
            elif actual and not detected:
                result_type = "fn"  # False negative - missed hallucination

        if not result_type:
            logger.debug("RL: grounding_threshold outcome missing result info, skipping")
            return

        # Build context key
        context_key = self._build_context_key(provider, response_type)

        # Compute reward
        reward = self._compute_reward(result_type)

        # Update Beta parameters (Thompson Sampling update)
        self._update_beta_params(context_key, threshold_used, result_type)

        # Update provider stats
        self._update_provider_stats(provider, result_type)

        # Save to database
        self._save_to_db(context_key, threshold_used, result_type, reward)

        self._total_decisions += 1

        logger.debug(
            f"RL: Grounding threshold for {provider}/{response_type} "
            f"threshold={threshold_used:.2f} result={result_type} reward={reward:.2f}"
        )

    def _build_context_key(self, provider: str, response_type: str) -> str:
        """Build context key for state lookup."""
        # Normalize response type
        if response_type not in self.RESPONSE_TYPES:
            response_type = "general"
        return f"{provider}:{response_type}"

    def _compute_reward(self, result_type: str) -> float:
        """Compute reward from verification result.

        Args:
            result_type: tp, tn, fp, or fn

        Returns:
            Reward value
        """
        rewards = {
            "tp": 0.1,  # True positive: correctly detected hallucination
            "tn": 0.1,  # True negative: correctly passed good response
            "fp": -1.0,  # False positive: wrongly flagged good response
            "fn": -2.0,  # False negative: missed hallucination (worst)
        }
        return rewards.get(result_type, 0.0)

    def _update_beta_params(self, context_key: str, threshold: float, result_type: str) -> None:
        """Update Beta distribution parameters for Thompson Sampling.

        For each threshold level, we track:
        - Alpha: proportional to successes (correct decisions at this threshold)
        - Beta: proportional to failures (incorrect decisions)

        Args:
            context_key: Provider + response type key
            threshold: Threshold value used
            result_type: Result of the decision
        """
        # Ensure context exists
        if context_key not in self._beta_params:
            self._beta_params[context_key] = {}
            for t in self.THRESHOLD_LEVELS:
                self._beta_params[context_key][t] = (self.PRIOR_ALPHA, self.PRIOR_BETA)

        # Find closest threshold level
        closest_threshold = min(self.THRESHOLD_LEVELS, key=lambda t: abs(t - threshold))

        # Get current parameters
        alpha, beta = self._beta_params[context_key][closest_threshold]

        # Update based on result
        if result_type in ("tp", "tn"):
            # Success - increase alpha
            alpha += self.learning_rate
        else:
            # Failure - increase beta
            beta += self.learning_rate

        self._beta_params[context_key][closest_threshold] = (alpha, beta)

    def _update_provider_stats(self, provider: str, result_type: str) -> None:
        """Update provider-level statistics."""
        cursor = self.db.cursor()

        # Ensure row exists
        cursor.execute(
            """
            INSERT OR IGNORE INTO grounding_threshold_stats
            (provider, true_positives, true_negatives, false_positives, false_negatives, last_updated)
            VALUES (?, 0, 0, 0, 0, ?)
            """,
            (provider, datetime.now().isoformat()),
        )

        # Update the appropriate counter
        column_map = {
            "tp": "true_positives",
            "tn": "true_negatives",
            "fp": "false_positives",
            "fn": "false_negatives",
        }
        column = column_map.get(result_type)
        if column:
            cursor.execute(
                f"""
                UPDATE grounding_threshold_stats
                SET {column} = {column} + 1, last_updated = ?
                WHERE provider = ?
                """,
                (datetime.now().isoformat(), provider),
            )

        self.db.commit()

    def _save_to_db(
        self, context_key: str, threshold: float, result_type: str, reward: float
    ) -> None:
        """Save parameters and history to database."""
        cursor = self.db.cursor()
        timestamp = datetime.now().isoformat()

        # Find closest threshold level
        closest_threshold = min(self.THRESHOLD_LEVELS, key=lambda t: abs(t - threshold))

        # Save Beta parameters
        alpha, beta = self._beta_params[context_key][closest_threshold]
        cursor.execute(
            """
            INSERT OR REPLACE INTO grounding_threshold_params
            (context_key, threshold, alpha, beta, sample_count, last_updated)
            VALUES (?, ?, ?, ?, COALESCE(
                (SELECT sample_count + 1 FROM grounding_threshold_params
                 WHERE context_key = ? AND threshold = ?), 1
            ), ?)
            """,
            (
                context_key,
                closest_threshold,
                alpha,
                beta,
                context_key,
                closest_threshold,
                timestamp,
            ),
        )

        # Save history
        cursor.execute(
            """
            INSERT INTO grounding_threshold_history
            (context_key, threshold, result_type, reward, timestamp)
            VALUES (?, ?, ?, ?, ?)
            """,
            (context_key, threshold, result_type, reward, timestamp),
        )

        self.db.commit()

    def get_recommendation(
        self, provider: str, model: str, task_type: str
    ) -> Optional[RLRecommendation]:
        """Get recommended threshold using Thompson Sampling.

        Uses Beta distribution sampling for explore/exploit balance.

        Args:
            provider: Provider name
            model: Model name (not used directly, but part of context)
            task_type: Response type / task type

        Returns:
            Recommendation with optimal threshold
        """
        context_key = self._build_context_key(provider, task_type)

        # Check if we have data for this context
        if context_key not in self._beta_params:
            return RLRecommendation(
                value=self.DEFAULT_THRESHOLD,
                confidence=0.3,
                reason="No learned data, using default threshold",
                sample_size=0,
                is_baseline=True,
            )

        # Thompson Sampling: sample from Beta distribution for each threshold
        samples = {}
        total_samples = 0
        for threshold, (alpha, beta) in self._beta_params[context_key].items():
            # Sample from Beta(alpha, beta)
            sample = random.betavariate(alpha, beta)
            samples[threshold] = sample
            total_samples += int(alpha + beta - 2 * self.PRIOR_ALPHA)

        # Select threshold with highest sampled value
        best_threshold = max(samples.keys(), key=lambda t: samples[t])
        best_sample = samples[best_threshold]

        # Compute confidence based on sample size and parameter certainty
        alpha, beta = self._beta_params[context_key][best_threshold]
        sample_count = int(alpha + beta - 2 * self.PRIOR_ALPHA)

        if sample_count < self.MIN_SAMPLES_FOR_CONFIDENCE:
            confidence = 0.3 + 0.2 * (sample_count / self.MIN_SAMPLES_FOR_CONFIDENCE)
            is_baseline = True
        else:
            # Confidence based on Beta distribution concentration
            variance = (alpha * beta) / ((alpha + beta) ** 2 * (alpha + beta + 1))
            confidence = min(0.95, 0.5 + 0.45 * (1 - math.sqrt(variance)))
            is_baseline = False

        return RLRecommendation(
            value=best_threshold,
            confidence=confidence,
            reason=f"Thompson sample={best_sample:.3f}, α={alpha:.1f}, β={beta:.1f}",
            sample_size=total_samples,
            is_baseline=is_baseline,
        )

    def get_optimal_threshold(
        self, provider: str, response_type: str = "general"
    ) -> Tuple[float, float]:
        """Get optimal threshold for a provider/response type.

        Convenience method that returns threshold and confidence.

        Args:
            provider: Provider name
            response_type: Type of response

        Returns:
            Tuple of (threshold, confidence)
        """
        rec = self.get_recommendation(provider, "", response_type)
        return (rec.value if rec else self.DEFAULT_THRESHOLD, rec.confidence if rec else 0.3)

    def get_provider_error_rates(self, provider: str) -> Dict[str, float]:
        """Get error rates for a provider.

        Args:
            provider: Provider name

        Returns:
            Dictionary with fp_rate, fn_rate, precision, recall
        """
        cursor = self.db.cursor()
        cursor.execute("SELECT * FROM grounding_threshold_stats WHERE provider = ?", (provider,))
        row = cursor.fetchone()

        if not row:
            return {"fp_rate": 0.0, "fn_rate": 0.0, "precision": 0.0, "recall": 0.0}

        row_dict = dict(row)
        tp = row_dict["true_positives"]
        tn = row_dict["true_negatives"]
        fp = row_dict["false_positives"]
        fn = row_dict["false_negatives"]

        total = tp + tn + fp + fn
        if total == 0:
            return {"fp_rate": 0.0, "fn_rate": 0.0, "precision": 0.0, "recall": 0.0}

        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0

        return {
            "fp_rate": fp / total,
            "fn_rate": fn / total,
            "precision": precision,
            "recall": recall,
            "total_samples": total,
        }

    def get_all_provider_stats(self) -> Dict[str, Dict[str, float]]:
        """Get error rates for all providers.

        Returns:
            Dictionary mapping provider to error rates
        """
        cursor = self.db.cursor()
        cursor.execute("SELECT DISTINCT provider FROM grounding_threshold_stats")
        providers = [row[0] for row in cursor.fetchall()]

        return {provider: self.get_provider_error_rates(provider) for provider in providers}

    def export_metrics(self) -> Dict[str, Any]:
        """Export learner metrics for monitoring.

        Returns:
            Dictionary with learner stats
        """
        # Compute average threshold per context
        avg_thresholds = {}
        for context_key, params in self._beta_params.items():
            # Weight thresholds by sample count
            weighted_sum = 0.0
            total_weight = 0.0
            for threshold, (alpha, beta) in params.items():
                weight = alpha + beta - 2 * self.PRIOR_ALPHA
                weighted_sum += threshold * weight
                total_weight += weight
            if total_weight > 0:
                avg_thresholds[context_key] = weighted_sum / total_weight

        return {
            "learner": self.name,
            "total_contexts": len(self._beta_params),
            "total_decisions": self._total_decisions,
            "threshold_levels": self.THRESHOLD_LEVELS,
            "default_threshold": self.DEFAULT_THRESHOLD,
            "learned_thresholds": avg_thresholds,
            "provider_stats": self.get_all_provider_stats(),
        }
