# ai_mouse/server/routes_collect.py
"""Collection routes: start, trace, skip."""
from __future__ import annotations

import json
import logging

from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel

from tools.collector import Collector

from .deps import SessionState, get_data_dir, get_state

logger = logging.getLogger(__name__)

router = APIRouter()

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


class CollectStartRequest(BaseModel):
    count: int = 100
    dist_min: int = 50
    dist_max: int = 800


class TraceRequest(BaseModel):
    meta: dict
    events: list[dict]


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


@router.post("/collect/start")
def collect_start(
    req: CollectStartRequest,
    state: SessionState = Depends(get_state),
) -> dict:
    traces_path = get_data_dir() / "traces.jsonl"
    collector = Collector(
        count=req.count,
        dist_min=req.dist_min,
        dist_max=req.dist_max,
        output_path=traces_path,
    )
    state.collector = collector
    return {"a": list(collector.a_pos), "b": list(collector.b_pos)}


@router.post("/collect/trace")
def collect_trace(
    trace: TraceRequest,
    state: SessionState = Depends(get_state),
) -> dict:
    if state.collector is None:
        raise HTTPException(
            status_code=400,
            detail="Collector not started. Call /api/collect/start first.",
        )

    traces_path = get_data_dir() / "traces.jsonl"
    traces_path.parent.mkdir(parents=True, exist_ok=True)
    record = {"meta": trace.meta, "events": trace.events}
    with traces_path.open("a", encoding="utf-8") as f:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")

    collector = state.collector
    collector.collected += 1
    remaining = collector.count - collector.collected

    if remaining > 0:
        collector.a_pos, collector.b_pos = collector._new_ab()
        return {
            "collected": collector.collected,
            "remaining": remaining,
            "a": list(collector.a_pos),
            "b": list(collector.b_pos),
        }
    else:
        return {
            "collected": collector.collected,
            "remaining": 0,
            "a": None,
            "b": None,
        }


@router.post("/collect/skip")
def collect_skip(
    state: SessionState = Depends(get_state),
) -> dict:
    if state.collector is None:
        raise HTTPException(
            status_code=400,
            detail="Collector not started. Call /api/collect/start first.",
        )

    collector = state.collector
    collector.a_pos, collector.b_pos = collector._new_ab()
    return {"a": list(collector.a_pos), "b": list(collector.b_pos)}