"""Implementation for the `chat_eval cli tune` command."""
from __future__ import annotations

import json
from datetime import datetime, timedelta, timezone


def build_audit_entry(
    *,
    step_id: int,
    ts: str,
    outcome: str,
    identity: dict,
    window_start_ts: str,
    window_end_ts: str,
    cleanup_dropped_doc_ids: list[str],
    **extra,
) -> dict:
    return {
        "step_id": step_id,
        "ts": ts,
        "outcome": outcome,
        **identity,
        "window_start_ts": window_start_ts,
        "window_end_ts": window_end_ts,
        "cleanup_dropped_doc_ids": cleanup_dropped_doc_ids,
        **extra,
    }


def build_identity_fields(turns_path_or_list) -> dict:
    from scripts.lib.chat_eval import judge as _judge
    from scripts.lib.chat_eval.rubric import RUBRIC_VERSION
    from scripts.lib.chat_eval.turn_log import iter_turns_compat

    fields = {
        "rubric_version": int(RUBRIC_VERSION),
        "judge_model": _judge._JUDGE_MODEL,
        "chat_model": None,
        "embedding_model": None,
        "catalog_artifact_version_id": None,
    }
    latest_judged_ts = ""
    for t in iter_turns_compat(turns_path_or_list):
        if not t.get("judge"):
            continue
        ts = t.get("ts", "")
        if ts > latest_judged_ts:
            latest_judged_ts = ts
            fields["chat_model"] = t.get("chat_model")
            fields["embedding_model"] = t.get("embedding_model")
            fields["catalog_artifact_version_id"] = t.get("catalog_artifact_version_id")
    return fields


def recent_util_history(turns_path_or_list, doc_id: str, *, end_ts: str,
                        days: int = 14) -> list[dict]:
    from scripts.lib.chat_eval.turn_log import iter_turns_compat

    end_dt = datetime.strptime(end_ts, "%Y-%m-%dT%H:%M:%SZ").replace(
        tzinfo=timezone.utc
    )
    out: list[dict] = []
    for offset in range(days - 1, -1, -1):
        bin_start = end_dt - timedelta(days=offset + 1)
        bin_end = end_dt - timedelta(days=offset)
        bin_start_iso = bin_start.strftime("%Y-%m-%dT%H:%M:%SZ")
        bin_end_iso = bin_end.strftime("%Y-%m-%dT%H:%M:%SZ")
        retrieved_count = 0
        used_count = 0
        for turn in iter_turns_compat(turns_path_or_list, since=bin_start_iso):
            if turn.get("ts", "") >= bin_end_iso or not turn.get("judge"):
                continue
            doc_ids_retrieved = {r["doc_id"] for r in (turn.get("retrieved") or [])}
            if doc_id in doc_ids_retrieved:
                retrieved_count += 1
            used = (turn.get("judge") or {}).get("docs_actually_used") or []
            if any(u.get("doc_id") == doc_id for u in used):
                used_count += 1
        if retrieved_count == 0:
            continue
        out.append({
            "bin_start_date": bin_start.date().isoformat(),
            "retrieved": retrieved_count,
            "used": used_count,
            "util_pct": round(100.0 * used_count / retrieved_count, 1),
        })
    return out


def _append_jsonl(path, entry: dict) -> None:
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(entry) + "\n")


def _load_previous_accepted(log_path) -> dict | None:
    previous = None
    if not log_path.is_file():
        return previous
    for line in log_path.read_text("utf-8").splitlines():
        if not line.strip():
            continue
        try:
            entry = json.loads(line)
        except json.JSONDecodeError:
            continue
        if entry.get("outcome") == "accepted":
            previous = entry
    return previous


def _cleanup_stale_penalties(*, state: dict, turns_path, now_ts: str, dry_run: bool) -> list[str]:
    cleanup_dropped: list[str] = []
    if not state.get("penalties"):
        return cleanup_dropped

    catalog_path = turns_path.parent / "catalog.jsonl"
    if not catalog_path.is_file():
        return cleanup_dropped

    known_ids: set[str] = set()
    for line in catalog_path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        try:
            rec = json.loads(line)
        except json.JSONDecodeError:
            continue
        doc_id = rec.get("doc_id")
        if doc_id:
            known_ids.add(doc_id)

    for doc_id in list(state["penalties"].keys()):
        if doc_id not in known_ids:
            state["penalties"].pop(doc_id)
            cleanup_dropped.append(doc_id)

    if cleanup_dropped and not dry_run:
        state["updated_at"] = now_ts
    return cleanup_dropped


def _write_skip(
    *,
    args,
    state: dict,
    now_ts: str,
    outcome: str,
    identity: dict,
    window_start_ts: str,
    window_end_ts: str,
    cleanup_dropped: list[str],
    **extra,
) -> None:
    entry = build_audit_entry(
        step_id=state["step_count"] + 1,
        ts=now_ts,
        outcome=outcome,
        identity=identity,
        window_start_ts=window_start_ts,
        window_end_ts=window_end_ts,
        cleanup_dropped_doc_ids=cleanup_dropped,
        **extra,
    )
    if not args.dry_run:
        _append_jsonl(args.tuning_log_path, entry)


def run(args) -> int:
    from scripts.lib.chat_eval import analytics, tuning, turn_log

    state_path = args.tuning_state_path
    log_path = args.tuning_log_path
    turns_path = args.log_path
    all_turns = [
        t for t in turn_log.load_all(turns_path)
        if turn_log.is_student_mode_turn(t)
    ]

    state = (
        json.loads(state_path.read_text("utf-8"))
        if state_path.is_file()
        else tuning.empty_state()
    )

    if args.unfreeze:
        if state["frozen"]:
            state["frozen"] = False
            state["frozen_reason"] = None
            state["updated_at"] = tuning.now_utc_iso()
            state_path.write_text(json.dumps(state), encoding="utf-8")
            print(f"Unfroze tuning state at {state_path}.")
        else:
            print("State already not frozen; no-op.")
        return 0

    now_ts = tuning.now_utc_iso()
    now_dt = datetime.strptime(now_ts, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc)
    window_end_ts = now_ts
    window_start_ts = (now_dt - timedelta(days=7)).strftime("%Y-%m-%dT%H:%M:%SZ")

    cleanup_dropped = _cleanup_stale_penalties(
        state=state, turns_path=turns_path, now_ts=now_ts, dry_run=args.dry_run
    )
    if cleanup_dropped and not args.dry_run:
        state_path.write_text(json.dumps(state), encoding="utf-8")

    identity = build_identity_fields(all_turns)

    if state["frozen"]:
        _write_skip(
            args=args,
            state=state,
            now_ts=now_ts,
            outcome="skipped_frozen",
            identity=identity,
            window_start_ts=window_start_ts,
            window_end_ts=window_end_ts,
            cleanup_dropped=cleanup_dropped,
            frozen_reason=state.get("frozen_reason"),
        )
        print(f"Skipped: state is frozen ({state.get('frozen_reason')!r}).")
        return 0

    prev_accepted = _load_previous_accepted(log_path)
    before_scores = None
    after_scores = None
    before_window_start_ts = None
    before_window_end_ts = None
    regression_check = None

    if prev_accepted:
        before_window_start_ts = prev_accepted["window_start_ts"]
        before_window_end_ts = prev_accepted["window_end_ts"]
        before_stats = analytics.score_summary_with_stats(
            all_turns, since=before_window_start_ts, until=before_window_end_ts
        )
        after_stats = analytics.score_summary_with_stats(
            all_turns, since=prev_accepted["ts"], until=window_end_ts
        )
        before_scores = {k: v["mean"] for k, v in before_stats.items()}
        after_scores = {k: v["mean"] for k, v in after_stats.items()}
        after_count = max((v["n"] for v in after_stats.values()), default=0)
        regression_check = tuning.check_regression(
            before_scores,
            after_scores,
            judged_turns=after_count,
            before_stats=before_stats,
            after_stats=after_stats,
        )

        if regression_check.get("skipped"):
            _write_skip(
                args=args,
                state=state,
                now_ts=now_ts,
                outcome="skipped_insufficient_turns",
                identity=identity,
                window_start_ts=window_start_ts,
                window_end_ts=window_end_ts,
                cleanup_dropped=cleanup_dropped,
                judged_turns_in_window=after_count,
                previous_step_id=prev_accepted["step_id"],
            )
            print(
                f"Skipped: only {after_count} judged turns since step "
                f"{prev_accepted['step_id']}; need >="
                f"{tuning.MIN_TURNS_FOR_REGRESSION_CHECK}."
            )
            return 0

        if regression_check["any_dimension_regressed"]:
            reason = (
                f"{regression_check['max_drop_dimension']} regressed "
                f"delta={regression_check['max_drop_value']} "
                f"after step {prev_accepted['step_id']}"
            )
            new_state = tuning.rollback_step(state, prev_accepted, reason=reason)
            rb_step_action = dict(prev_accepted.get("step_action") or {})
            if rb_step_action.get("doc_id"):
                rb_step_action["recent_util_history"] = recent_util_history(
                    all_turns, rb_step_action["doc_id"], end_ts=window_end_ts
                )
            rb_entry = build_audit_entry(
                step_id=new_state["step_count"] + 1,
                ts=now_ts,
                outcome="rolled_back",
                identity=identity,
                window_start_ts=window_start_ts,
                window_end_ts=window_end_ts,
                cleanup_dropped_doc_ids=cleanup_dropped,
                before_window_start_ts=before_window_start_ts,
                before_window_end_ts=before_window_end_ts,
                rolled_back_step_id=prev_accepted["step_id"],
                regression_check=regression_check,
                scores_before=before_scores,
                scores_after=after_scores,
                step_action=rb_step_action,
            )
            if not args.dry_run:
                state_path.write_text(json.dumps(new_state), encoding="utf-8")
                _append_jsonl(log_path, rb_entry)
            print(f"Regression detected: {reason}. Rolled back + frozen.")
            return 0

    util = analytics.doc_utilization(all_turns)
    step_action = tuning.propose_step(util, current_state=state)

    if step_action is None:
        _write_skip(
            args=args,
            state=state,
            now_ts=now_ts,
            outcome="skipped_no_eligible_doc",
            identity=identity,
            window_start_ts=window_start_ts,
            window_end_ts=window_end_ts,
            cleanup_dropped=cleanup_dropped,
            judged_turns_seen=sum(1 for t in all_turns if t.get("judge")),
        )
        print("No eligible doc for a step; skipped.")
        return 0

    new_state = tuning.apply_step(state, step_action)
    step_action["recent_util_history"] = recent_util_history(
        all_turns, step_action["doc_id"], end_ts=window_end_ts
    )

    if prev_accepted is None:
        cold_after_stats = analytics.score_summary_with_stats(
            all_turns, since=window_start_ts, until=window_end_ts
        )
        cold_after_scores = {k: v["mean"] for k, v in cold_after_stats.items()}
        after_scores = cold_after_scores
        regression_check = {"skipped": True, "reason": "cold_start"}

    entry = build_audit_entry(
        step_id=new_state["step_count"],
        ts=now_ts,
        outcome="accepted",
        identity=identity,
        window_start_ts=window_start_ts,
        window_end_ts=window_end_ts,
        cleanup_dropped_doc_ids=cleanup_dropped,
        before_window_start_ts=before_window_start_ts,
        before_window_end_ts=before_window_end_ts,
        scores_before=before_scores,
        scores_after=after_scores,
        regression_check=regression_check,
        step_action=step_action,
    )

    if args.dry_run:
        print(json.dumps({"would_apply": step_action, "audit_entry": entry}, indent=2))
        return 0

    state_path.write_text(json.dumps(new_state), encoding="utf-8")
    _append_jsonl(log_path, entry)
    print(
        f"Applied step #{new_state['step_count']} on doc "
        f"{step_action['doc_id']}: {step_action['penalty_before']} -> "
        f"{step_action['penalty_after']}"
    )
    return 0