# ai_mouse/server/routes_train.py
"""Training and status routes."""
from __future__ import annotations

import asyncio
import json
import logging
from pathlib import Path
from typing import AsyncGenerator

from fastapi import APIRouter
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from .deps import get_data_dir

logger = logging.getLogger(__name__)

router = APIRouter()

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _paths() -> tuple[Path, Path, Path]:
    data_dir = get_data_dir()
    return (
        data_dir / "traces.jsonl",
        data_dir / "models_v2",
        data_dir / "models_v2_pretrained",
    )


def _trace_count() -> int:
    traces_path, _, _ = _paths()
    if not traces_path.exists():
        return 0
    return sum(
        1
        for line in traces_path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    )


def _model_trained() -> bool:
    _, models_dir, _ = _paths()
    return (models_dir / "flow_model.pt").exists()


# ---------------------------------------------------------------------------
# Request models
# ---------------------------------------------------------------------------


class TrainRequest(BaseModel):
    epochs: int = 200
    data_path: str | None = None
    output_dir: str | None = None


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------


@router.get("/status")
def get_status() -> dict:
    return {"trace_count": _trace_count(), "model_trained": _model_trained()}


async def _train_sse_generator(req: TrainRequest) -> AsyncGenerator[str, None]:
    """Run training in a thread via asyncio.to_thread, yield SSE events via asyncio.Queue."""
    queue: asyncio.Queue[dict] = asyncio.Queue()

    def callback(msg: dict) -> None:
        queue.put_nowait(msg)

    async def run_training_async() -> None:
        from tools.trainer import train

        traces_path, models_dir, pretrained_dir = _paths()
        data_path = Path(req.data_path) if req.data_path else traces_path
        output_dir = Path(req.output_dir) if req.output_dir else models_dir

        # Auto-detect pretrained checkpoint and switch to fine-tune mode
        resume_from: Path | None = None
        effective_lr = 3e-4
        if (pretrained_dir / "flow_model.pt").exists():
            resume_from = pretrained_dir
            effective_lr = 1e-5  # fine-tune lr
            logger.info("Detected pretrained checkpoint, fine-tuning at lr=%g", effective_lr)
            queue.put_nowait({
                "info": f"Detected pretrained checkpoint at {pretrained_dir.name}, "
                        f"running fine-tune at lr={effective_lr}",
            })

        try:
            await asyncio.to_thread(
                train,
                data_path=data_path,
                output_dir=output_dir,
                epochs=req.epochs,
                lr=effective_lr,
                progress_callback=callback,
                resume_from=resume_from,
            )
        except Exception as exc:  # noqa: BLE001
            queue.put_nowait({"error": str(exc)})

    task = asyncio.create_task(run_training_async())

    while True:
        msg = await queue.get()
        yield f"data: {json.dumps(msg)}\n\n"
        if msg.get("done") or msg.get("error"):
            break

    await task


@router.post("/train")
async def train_model(req: TrainRequest) -> StreamingResponse:
    return StreamingResponse(
        _train_sse_generator(req),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )