from __future__ import annotations

from datetime import datetime, timedelta
from typing import Dict, Any, Optional

import duckdb

from ..intelligence.decision_logger import DecisionLogger
from ..intelligence.bandit_repo import BanditRepo
from ..intelligence.policy_store import PolicyStore


CLAMP = lambda v, lo, hi: max(lo, min(hi, v))


class LearningPipeline:
    def __init__(
        self,
        decision_logger: DecisionLogger,
        bandit_repo: BanditRepo,
        policy_store: PolicyStore,
    ):
        self.logger = decision_logger
        self.bandits = bandit_repo
        self.policies = policy_store

    async def update_from_feedback(
        self, *, account_id: str, segment: Optional[str] = None, lookback_days: int = 30
    ) -> Dict[str, Any]:
        seg = segment or "default"

        # 1) Read feedback events and map to rewards per arm
        start_ts = datetime.utcnow() - timedelta(days=int(lookback_days))
        accepted = 0
        reverted = 0
        totals = 0
        per_arm: Dict[str, Dict[str, int]] = {}

        with duckdb.connect(str(self.logger.duckdb_path)) as conn:
            rows = conn.execute(
                """
                SELECT d.threshold_policy as arm,
                       f.outcome as outcome
                FROM feedback_events f
                JOIN decision_log d ON d.job_id = f.job_id
                WHERE d.tenant_id = ?
                  AND d.segment = ?
                  AND f.ts > ?
                """,
                [str(account_id), str(seg), start_ts],
            ).fetchall()

        for arm, outcome in rows:
            totals += 1
            if outcome == "accepted":
                accepted += 1
                delta_w, delta_t = 1, 1
            elif outcome == "reverted":
                reverted += 1
                delta_w, delta_t = 0, 1
            else:
                # treat other outcomes neutral as a trial
                delta_w, delta_t = 0, 1
            per = per_arm.setdefault(arm or "default", {"wins": 0, "trials": 0})
            per["wins"] += delta_w
            per["trials"] += delta_t

        # 2) Update bandit arm stats
        for arm, stats in per.items():
            self.bandits.upsert_arm_stats(
                account_id,
                seg,
                arm,
                wins_delta=stats["wins"],
                trials_delta=stats["trials"],
            )

        snapshot = self.bandits.get_ucb_snapshot(account_id, seg)
        chosen_arm = snapshot["arms"][0]["name"] if snapshot["arms"] else "default"

        # 3) Adjust thresholds with small nudges based on acceptance rate
        policy = await self.policies.get_policy(account_id, seg)
        suggest = float(policy.get("suggest", 0.70))
        min_gap = float(policy.get("min_gap", 0.10))
        auto_link = float(policy.get("auto_link", suggest + min_gap))

        # Acceptance rate proxy
        acc_rate = (accepted / totals) if totals > 0 else None
        delta = 0.0
        if acc_rate is not None:
            if acc_rate < 0.5:
                delta = -0.01
            elif acc_rate > 0.7:
                delta = +0.01

        new_suggest = CLAMP(suggest + delta, 0.50, 0.85)
        new_min_gap = CLAMP(min_gap, 0.05, 0.20)
        new_auto_link = CLAMP(new_suggest + new_min_gap, 0.55, 0.99)

        new_policy = {
            "suggest": round(new_suggest, 4),
            "min_gap": round(new_min_gap, 4),
            "auto_link": round(new_auto_link, 4),
            # weights influenced by chosen arm (optional small tweaks)
            "weights": self._weights_for_arm(
                chosen_arm, base=policy.get("weights") or {}
            ),
            "version": int(policy.get("version") or 1) + 1,
        }

        await self.policies.set_policy(account_id, seg, new_policy)

        return {
            "account_id": account_id,
            "segment": seg,
            "lookback_days": lookback_days,
            "feedback": {
                "accepted": accepted,
                "reverted": reverted,
                "total": totals,
                "accept_rate": acc_rate,
            },
            "bandit": snapshot,
            "chosen_arm": chosen_arm,
            "policy_before": {
                "suggest": suggest,
                "auto_link": auto_link,
                "min_gap": min_gap,
            },
            "policy_after": new_policy,
        }

    def _weights_for_arm(self, arm: str, base: Dict[str, float]) -> Dict[str, float]:
        # conservative bounded adjustments
        fa = float(base.get("fuzzy_a", 1.0))
        fb = float(base.get("fuzzy_b", 1.0))
        ex = float(base.get("exact", 1.0))
        if not arm:
            return {"fuzzy_a": fa, "fuzzy_b": fb, "exact": ex}
        a = arm.lower()
        if "precision" in a:
            ex = CLAMP(ex + 0.05, 0.5, 2.0)
            fb = CLAMP(fb - 0.05, 0.5, 2.0)
        elif "recall" in a:
            fb = CLAMP(fb + 0.05, 0.5, 2.0)
            ex = CLAMP(ex - 0.05, 0.5, 2.0)
        return {"fuzzy_a": round(fa, 4), "fuzzy_b": round(fb, 4), "exact": round(ex, 4)}
