from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Iterable

from prompt_auto_lab.orchestrator.state_store import SharedState
from prompt_auto_lab.orchestrator.workflow import WorkflowConfig


def _extract_fields_from_workflow(workflow: WorkflowConfig | None) -> tuple[set[str], set[str]]:
    """
    Return (iteration_fields, top_level_fields) derived from workflow.state_fields.
    iteration_fields exclude the leading iteration[]/iteration. prefix.
    """
    iteration_fields: set[str] = set()
    top_fields: set[str] = set()
    if workflow is None:
        return iteration_fields, top_fields

    for f in workflow.state_fields:
        key = f.key
        if key.startswith("iteration[]."):
            iteration_fields.add(key.removeprefix("iteration[]."))
        elif key.startswith("iteration."):
            iteration_fields.add(key.removeprefix("iteration."))
        else:
            top_fields.add(key)
    return iteration_fields, top_fields


class HistoryUtils:
    """Helpers for summarizing current node state and reading persisted history."""

    def __init__(
        self,
        shared_state: SharedState,
        workflow: WorkflowConfig | None = None,
        versions_path: Path | None = None,
    ) -> None:
        self._shared_state = shared_state
        self._workflow = workflow
        self._versions_path = versions_path
        self._iter_fields, self._top_fields = _extract_fields_from_workflow(workflow)

    async def current_snapshot(self) -> dict[str, Any]:
        """
        Return a concise view of the active node and selected top-level fields.
        Fields are chosen from workflow.state_fields when provided; falls back to common DAG fields.
        """
        snap = await self._shared_state.read(None)
        iteration = snap.get("iteration") or {}
        idx = snap.get("iteration_idx") or "root"
        node = iteration.get(str(idx)) or iteration.get(idx) or {}

        # # Structural fields always included
        # result: dict[str, Any] = {
        #     "node_id": idx,
        #     "parent_node": node.get("parent_node"),
        #     "analysis_source_node": node.get("analysis_source_node"),
        # }

        # # Iteration-scoped fields from workflow (or common defaults)
        # iter_fields = self._iter_fields or {
        #     "experiment_id",
        #     "experiment_link",
        #     "metrics",
        #     "overall_score",
        #     "analysis_results_path",
        #     "current_prompt_file_path",
        #     "if_optimized",
        #     "continue",
        #     "reason",
        #     "commit_id"
        # }
        # for key in iter_fields:
        #     result[key] = node.get(key)
        result= node.copy()
        # Top-level fields from workflow (or common defaults)
        top_fields = self._top_fields or {"continue", "stop_reason", "max_rounds", "branch", "version"}
        for key in top_fields:
            result[f"top.{key}"] = snap.get(key)

        return result

    def list_history(self, limit: int = 20) -> list[dict[str, Any]]:
        """Return recent persisted versions (if any) from versions.jsonl."""
        records = list(self._iter_versions())
        if limit > 0:
            records = records[-limit:]
        return records

    def best_score(self) -> dict[str, Any] | None:
        """Return the record with the best overall_score, if available."""
        best: dict[str, Any] | None = None
        for rec in self._iter_versions():
            score = self._safe_score(rec)
            if score is None:
                continue
            if best is None or score > (self._safe_score(best) or float("-inf")):
                best = rec
        return best

    def _iter_versions(self) -> Iterable[dict[str, Any]]:
        path = self._versions_path
        if path is None or not path.exists():
            return []
        out: list[dict[str, Any]] = []
        for line in path.read_text(encoding="utf-8").splitlines():
            if not line.strip():
                continue
            try:
                obj = json.loads(line)
                if isinstance(obj, dict):
                    out.append(obj)
            except Exception:
                continue
        return out

    @staticmethod
    def _safe_score(rec: dict[str, Any]) -> float | None:
        val = rec.get("overall_score")
        try:
            return float(val) if val is not None else None
        except Exception:
            return None
