"""The task -> result pipeline. This is the single integration point for clients.

Anything that wants to run episodes — the CLI, an SFT-data generator, an eval
harness, an RL training loop — constructs an EpisodePipeline and calls
run_episode(task, policy). Everything else (workspace, validation, queue, compute,
grading, storage, costs) happens behind this call.
"""

from __future__ import annotations

import uuid

import verifiers

from autonomous_trainer.agent.loop import AgentLimits, run_agent_episode
from autonomous_trainer.base_models import base_model_info
from autonomous_trainer.grading import GradeWeights, grade_episode
from autonomous_trainer.jobs import ComputeProvider, JobQueue, MockComputeProvider
from autonomous_trainer.jobs.eval_cache import EvalCache, cache_key
from autonomous_trainer.policy import Policy
from autonomous_trainer.schemas import (
    EpisodeResult,
    JobConfig,
    JobRecord,
    JobSpec,
    TaskSpec,
    utcnow,
)
from autonomous_trainer.settings import Settings, get_settings
from autonomous_trainer.store import RunStore
from autonomous_trainer.validation import validate_workspace
from autonomous_trainer.workspace import Workspace


def _load_job_config(workspace: Workspace) -> JobConfig:
    """Parse the (already validated) job.toml into a JobConfig."""
    import tomllib

    raw = tomllib.loads(workspace.read_file("job.toml"))
    return JobConfig(
        **{
            **raw.get("model", {}),
            **raw.get("training", {}),
            "rollout_max_turns": raw.get("rollout", {}).get("max_turns", 4),
        }
    )


class EpisodePipeline:
    def __init__(
        self,
        settings: Settings | None = None,
        provider: ComputeProvider | None = None,
        grade_weights: GradeWeights = GradeWeights(),
    ):
        self.settings = settings or get_settings()
        self.provider = provider or MockComputeProvider()
        self.grade_weights = grade_weights
        self.store = RunStore(self.settings.runs_dir)
        self.queue = JobQueue(self.settings.runs_dir / "jobs")
        self.eval_cache = EvalCache(self.settings.runs_dir / "eval_cache.json")
        # best_pre per eval env, snapshotted on first use and FROZEN for this
        # pipeline's lifetime: grading must use one constant per task across a
        # whole training run (a mid-run shift would give identical rollouts in
        # one GRPO group different rewards). Seed via scripts/seed_baselines.py
        # before RL; None (unseeded) falls back to absolute-post grading.
        self._best_pre: dict[str, float | None] = {}

    async def run_episode(self, task: TaskSpec, policy: Policy) -> EpisodeResult:
        s = self.settings
        run = self.store.new_run()
        run.log("episode_start", task_id=task.id, model=policy.model_id)
        workspace = Workspace.seed(run.workspace_dir, s.template_dir)

        async def validate():
            return await validate_workspace(workspace.root, task, timeout_s=s.validation_timeout_s)

        # Cached baseline scores per allowed base model, exposed to the agent via
        # the get_baseline_scores tool (empirical model selection — the cache
        # fills as the fleet runs jobs, so later episodes inherit measurements).
        baselines = {
            model: self.eval_cache.get(
                cache_key(model, task.eval_env_module, verifiers.__version__)
            )
            for model in task.constraints.allowed_base_models
        }

        outcome = await run_agent_episode(
            policy,
            task,
            workspace,
            validate,
            run,
            AgentLimits(
                max_turns=s.max_agent_turns,
                max_submission_attempts=s.max_submission_attempts,
                max_usd=s.max_episode_usd,
            ),
            baselines=baselines,
        )

        if s.record_trajectory:
            import json as _json

            (run.root / "messages.json").write_text(
                _json.dumps(outcome.messages, indent=2, default=str)
            )

        job: JobRecord | None = None
        if outcome.validated:
            config = _load_job_config(workspace)
            pre_key = cache_key(config.base_model, task.eval_env_module, verifiers.__version__)
            spec = JobSpec(
                job_id=f"job-{run.run_id}-{uuid.uuid4().hex[:6]}",
                run_id=run.run_id,
                task_id=task.id,
                workspace_path=str(workspace.root),
                config=config,
                verifiers_version=verifiers.__version__,
                eval_env_module=task.eval_env_module,
                eval_reward_pre_cached=self.eval_cache.get(pre_key),
                tool_call_parser=base_model_info(config.base_model).tool_call_parser,
            )
            job = await self.queue.submit_and_process(spec, self.provider)
            if (
                spec.eval_reward_pre_cached is None
                and job.result is not None
                and "eval_reward_pre" in job.result.metrics
            ):
                self.eval_cache.put(pre_key, job.result.metrics["eval_reward_pre"])
            run.log(
                "job_processed",
                job_id=spec.job_id,
                provider=self.provider.name,
                status=job.status.value,
                pre_eval_cached=spec.eval_reward_pre_cached is not None,
                result=job.result.model_dump() if job.result else None,
            )

        if task.eval_env_module not in self._best_pre:
            measured = [v for v in baselines.values() if v is not None]
            self._best_pre[task.eval_env_module] = max(measured) if measured else None
            if not measured:
                run.log("best_pre_missing", eval_env=task.eval_env_module)

        reward, components = grade_episode(
            outcome.validated,
            outcome.submission_attempts,
            s.max_submission_attempts,
            job,
            self.grade_weights,
            total_tokens=outcome.usage.prompt_tokens + outcome.usage.completion_tokens,
            token_budget=s.episode_token_budget,
            train_time_budget_s=task.train_time_budget_s,
            best_pre=self._best_pre[task.eval_env_module],
        )

        result = EpisodeResult(
            run_id=run.run_id,
            task_id=task.id,
            model=policy.model_id,
            reward=reward,
            reward_components=components,
            validated=outcome.validated,
            submission_attempts=outcome.submission_attempts,
            stop_reason=outcome.stop_reason,
            job=job,
            usage=outcome.usage,
            trajectory_path=str(run.trajectory_path),
            workspace_path=str(workspace.root),
            finished_at=utcnow(),
        )
        run.log("episode_end", reward=reward, components=components, usd=outcome.usage.usd)
        run.finalize(result)
        self.store.record_cost(result)
        return result