"""Learning feedback persistence: playbook stats and bounded policy adaptation.

This module provides a conservative cross-run loop:
- Update playbook success/failure counters per pattern key.
- Maintain bounded policy deltas per task type.
- Apply small, safe adjustments to runtime defaults on subsequent runs.
"""

from __future__ import annotations

import json
from datetime import datetime
from pathlib import Path

from backend.config import MEMORY_DIR


_PLAYBOOK_FILE = MEMORY_DIR / "playbook_stats.json"
_POLICY_FEEDBACK_FILE = MEMORY_DIR / "policy_feedback.json"


def _load_json(path: Path, fallback: dict) -> dict:
    try:
        if path.exists():
            data = json.loads(path.read_text(encoding="utf-8"))
            if isinstance(data, dict):
                return data
    except Exception:
        pass
    return dict(fallback)


def _save_json(path: Path, data: dict) -> None:
    try:
        MEMORY_DIR.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
    except Exception:
        pass


def derive_playbook_key(goal: str, task_state: dict | None = None) -> str:
    task = dict(task_state or {})
    task_type = str(task.get("task_type") or "generic").strip().lower() or "generic"
    provider_scope = str(task.get("provider_scope") or "default").strip().lower() or "default"
    time_filter = str(task.get("time_filter") or "none").strip().lower() or "none"
    return f"{task_type}|{provider_scope}|{time_filter}"


def update_playbook_stats(
    *,
    goal: str,
    task_state: dict | None,
    run_kpis: dict,
    hypothesis_summary: dict | None = None,
) -> dict:
    store = _load_json(_PLAYBOOK_FILE, {"version": 1, "patterns": {}})
    patterns = store.get("patterns") if isinstance(store.get("patterns"), dict) else {}
    key = derive_playbook_key(goal, task_state)
    row = patterns.get(key) if isinstance(patterns.get(key), dict) else {}

    attempts = int(row.get("attempts") or 0) + 1
    completed = bool(run_kpis.get("completed"))
    first_pass = bool(run_kpis.get("first_pass_success"))
    successes = int(row.get("successes") or 0) + (1 if completed else 0)
    first_pass_successes = int(row.get("first_pass_successes") or 0) + (1 if first_pass else 0)

    total_ms_prev = int(row.get("total_time_ms") or 0)
    total_ms = total_ms_prev + int(run_kpis.get("time_to_solution_ms") or 0)
    avg_time_ms = int(total_ms / max(1, attempts))

    row.update(
        {
            "attempts": attempts,
            "successes": successes,
            "first_pass_successes": first_pass_successes,
            "success_rate": round(successes / max(1, attempts), 4),
            "first_pass_rate": round(first_pass_successes / max(1, attempts), 4),
            "avg_time_ms": avg_time_ms,
            "total_time_ms": total_ms,
            "last_updated": datetime.now().isoformat(),
            "last_active_hypothesis": str((hypothesis_summary or {}).get("active_hypothesis_id") or ""),
            "last_confidence": float((hypothesis_summary or {}).get("active_confidence") or 0.0),
        }
    )

    if attempts >= 6 and row["success_rate"] >= 0.75 and row["first_pass_rate"] >= 0.55:
        row["status"] = "promoted"
    elif attempts >= 6 and row["success_rate"] < 0.45:
        row["status"] = "demoted"
    else:
        row["status"] = row.get("status") or "candidate"

    patterns[key] = row
    store["patterns"] = patterns
    _save_json(_PLAYBOOK_FILE, store)
    return {"key": key, **row}


def get_playbook_stats() -> dict:
    return _load_json(_PLAYBOOK_FILE, {"version": 1, "patterns": {}})


def apply_policy_feedback_adjustment(
    run_settings: dict,
    *,
    task_state: dict | None,
) -> dict:
    """Return bounded adjusted policy defaults for this run.

    Output keys mirror run settings policy keys.
    """
    store = _load_json(_POLICY_FEEDBACK_FILE, {"version": 1, "task_types": {}})
    task_types = store.get("task_types") if isinstance(store.get("task_types"), dict) else {}
    task_type = str((task_state or {}).get("task_type") or "generic").strip().lower() or "generic"
    profile = task_types.get(task_type) if isinstance(task_types.get(task_type), dict) else {}

    def _i(value: object, default: int) -> int:
        try:
            return int(value)
        except Exception:
            return int(default)

    base_steps = max(1, _i(run_settings.get("policy_max_steps_per_run", 8), 8))
    base_calls = max(1, _i(run_settings.get("policy_max_tool_calls_per_run", 20), 20))
    base_tokens = max(1, _i(run_settings.get("policy_max_total_tokens_per_run", 120000), 120000))
    base_runtime = max(10, _i(run_settings.get("policy_max_runtime_seconds", 900), 900))

    delta_steps = _i(profile.get("delta_steps", 0), 0)
    delta_calls = _i(profile.get("delta_tool_calls", 0), 0)
    delta_tokens = _i(profile.get("delta_tokens", 0), 0)
    delta_runtime = _i(profile.get("delta_runtime_seconds", 0), 0)

    # Hard safety bounds.
    return {
        "task_type": task_type,
        "policy_max_steps_per_run": min(30, max(3, base_steps + delta_steps)),
        "policy_max_tool_calls_per_run": min(80, max(8, base_calls + delta_calls)),
        "policy_max_total_tokens_per_run": min(500000, max(40000, base_tokens + delta_tokens)),
        "policy_max_runtime_seconds": min(3600, max(120, base_runtime + delta_runtime)),
        "profile": profile,
    }


def update_policy_feedback_from_kpis(task_state: dict | None, run_kpis: dict) -> dict:
    store = _load_json(_POLICY_FEEDBACK_FILE, {"version": 1, "task_types": {}})
    task_types = store.get("task_types") if isinstance(store.get("task_types"), dict) else {}
    task_type = str((task_state or {}).get("task_type") or "generic").strip().lower() or "generic"
    row = task_types.get(task_type) if isinstance(task_types.get(task_type), dict) else {}

    attempts = int(row.get("attempts") or 0) + 1
    row["attempts"] = attempts
    row["last_updated"] = datetime.now().isoformat()
    row["last_run"] = {
        "first_pass_success": bool(run_kpis.get("first_pass_success")),
        "recovery_count": int(run_kpis.get("recovery_count") or 0),
        "time_to_solution_ms": int(run_kpis.get("time_to_solution_ms") or 0),
        "tool_call_efficiency": float(run_kpis.get("tool_call_efficiency") or 0.0),
    }

    delta_steps = int(row.get("delta_steps") or 0)
    delta_calls = int(row.get("delta_tool_calls") or 0)
    delta_tokens = int(row.get("delta_tokens") or 0)
    delta_runtime = int(row.get("delta_runtime_seconds") or 0)

    first_pass_success = bool(run_kpis.get("first_pass_success"))
    recovery_count = int(run_kpis.get("recovery_count") or 0)
    runtime_ms = int(run_kpis.get("time_to_solution_ms") or 0)
    tool_eff = float(run_kpis.get("tool_call_efficiency") or 0.0)

    # EMA smoothing (alpha=0.3) to dampen oscillation.
    _EMA_ALPHA = 0.3
    if not first_pass_success and recovery_count >= 2:
        target_steps = min(8, delta_steps + 2)
        target_calls = min(24, delta_calls + 4)
        target_runtime = min(900, delta_runtime + 120)
    elif first_pass_success and tool_eff >= 0.8 and runtime_ms < 120000:
        target_steps = max(-3, delta_steps - 2)
        target_calls = max(-8, delta_calls - 2)
        target_runtime = delta_runtime
    else:
        target_steps = delta_steps
        target_calls = delta_calls
        target_runtime = delta_runtime

    delta_steps = round(delta_steps + _EMA_ALPHA * (target_steps - delta_steps))
    delta_calls = round(delta_calls + _EMA_ALPHA * (target_calls - delta_calls))
    delta_runtime = round(delta_runtime + _EMA_ALPHA * (target_runtime - delta_runtime))

    if runtime_ms > 300000:
        target_runtime_long = min(1200, delta_runtime + 120)
        delta_runtime = round(delta_runtime + _EMA_ALPHA * (target_runtime_long - delta_runtime))
        target_tokens_long = min(120000, delta_tokens + 10000)
        delta_tokens = round(delta_tokens + _EMA_ALPHA * (target_tokens_long - delta_tokens))

    row["delta_steps"] = int(delta_steps)
    row["delta_tool_calls"] = int(delta_calls)
    row["delta_tokens"] = int(delta_tokens)
    row["delta_runtime_seconds"] = int(delta_runtime)

    task_types[task_type] = row
    store["task_types"] = task_types
    _save_json(_POLICY_FEEDBACK_FILE, store)
    return {"task_type": task_type, **row}


def get_policy_feedback() -> dict:
    return _load_json(_POLICY_FEEDBACK_FILE, {"version": 1, "task_types": {}})
