from __future__ import annotations

import os
import threading
from typing import Any

import yaml
from fastapi import FastAPI, HTTPException, Request, Response

from argus.core.ids import new_task_id
from argus.ray.models import JobSpec, RayConfig

from .config import V2Config
from .db import Db
from .scheduler import Scheduler


def _utc_now_iso() -> str:
    from datetime import datetime

    return datetime.utcnow().replace(microsecond=0).isoformat() + "Z"


def _load_yaml_file(path: str) -> dict[str, Any]:
    with open(path, "r", encoding="utf-8") as f:
        obj = yaml.safe_load(f) or {}
    if not isinstance(obj, dict):
        raise ValueError("config yaml must be a mapping")
    return obj


def create_app(config_path: str) -> FastAPI:
    root = _load_yaml_file(config_path)
    ray_cfg = RayConfig.from_dict(root)
    v2_cfg = V2Config.from_root_dict(root)

    db = Db(v2_cfg.sqlite.db_path)
    db.init()

    scheduler = Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
    stop_flag = threading.Event()
    tool = scheduler.tool

    app = FastAPI(title="mvp-v2", version="2.0")

    def _auth(req: Request) -> dict[str, Any]:
        token_env = v2_cfg.auth.token_env
        admin_token = os.environ.get(token_env, "")
        if not admin_token:
            raise HTTPException(status_code=500, detail=f"missing token env: {token_env}")

        auth = req.headers.get("authorization") or ""
        if not auth.startswith("Bearer "):
            raise HTTPException(status_code=401, detail="missing bearer token")
        got = auth.removeprefix("Bearer ").strip()

        if got == admin_token:
            return {"user_id": "admin", "is_admin": True}

        info = db.resolve_token(got)
        if not info:
            raise HTTPException(status_code=401, detail="invalid token")
        if info.get("state") != "ACTIVE":
            raise HTTPException(status_code=403, detail="user disabled")
        return {"user_id": info["user_id"], "is_admin": False}

    def _require_admin(req: Request) -> dict[str, Any]:
        subject = _auth(req)
        if not subject.get("is_admin"):
            raise HTTPException(status_code=403, detail="admin required")
        return subject

    @app.on_event("startup")
    def _startup() -> None:
        t = threading.Thread(target=scheduler.run_forever, args=(stop_flag,), daemon=True)
        t.start()

    @app.on_event("shutdown")
    def _shutdown() -> None:
        stop_flag.set()

    @app.post("/api/v2/users")
    async def create_user(req: Request) -> dict[str, Any]:
        _require_admin(req)
        obj = await req.json()
        if not isinstance(obj, dict):
            raise HTTPException(status_code=400, detail="body must be a JSON object")
        user_id = str(obj.get("user_id") or "").strip()
        if not user_id:
            raise HTTPException(status_code=400, detail="missing user_id")
        display_name = obj.get("display_name")
        try:
            row = db.create_user(user_id=user_id, display_name=str(display_name) if display_name is not None else None)
        except Exception as e:
            raise HTTPException(status_code=409, detail=f"user create failed: {e!r}")
        return {"user_id": row.get("user_id", user_id), "state": row.get("state", "ACTIVE")}

    @app.post("/api/v2/users/{user_id}/tokens")
    async def issue_token(user_id: str, req: Request) -> dict[str, Any]:
        _require_admin(req)
        u = db.get_user(user_id)
        if not u:
            raise HTTPException(status_code=404, detail="user not found")
        token = db.issue_token(user_id=user_id)
        return {"user_id": user_id, "token": token}

    @app.post("/api/v2/users/{user_id}:disable")
    async def disable_user(user_id: str, req: Request) -> dict[str, Any]:
        _require_admin(req)
        u = db.get_user(user_id)
        if not u:
            raise HTTPException(status_code=404, detail="user not found")
        db.disable_user(user_id=user_id)
        return {"user_id": user_id, "state": "DISABLED"}

    @app.post("/api/v2/tasks")
    async def submit_task(req: Request) -> dict[str, Any]:
        subject = _auth(req)
        body = (await req.body()).decode("utf-8")
        obj = yaml.safe_load(body) or {}
        if not isinstance(obj, dict):
            raise HTTPException(status_code=400, detail="jobspec must be a YAML mapping")

        try:
            spec = JobSpec.from_dict(obj)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"invalid jobspec: {e!r}")

        # v2.5 constraint: training inputs must come from /private/common (dev/prod统一)。
        common_prefix = ray_cfg.shared_root.rstrip("/") + "/common/"
        for k, v in (("code_path", spec.code_path), ("train_file", spec.train_file), ("val_file", spec.val_file)):
            if v is None:
                continue
            if not str(v).startswith(common_prefix):
                raise HTTPException(status_code=400, detail=f"{k} must start with {common_prefix}")

        task_id = new_task_id(spec.workload, user_id=str(subject["user_id"]))
        db.create_task_v25(
            task_id=task_id,
            user_id=str(subject["user_id"]),
            workload=spec.workload,
            jobspec_yaml=body,
            nnodes=spec.nnodes,
            n_gpus_per_node=spec.n_gpus_per_node,
        )
        return {"task_id": task_id, "state": "QUEUED"}

    @app.get("/api/v2/tasks/{task_id}")
    async def get_task(task_id: str, req: Request) -> dict[str, Any]:
        subject = _auth(req)
        row = db.get_task(task_id)
        if not row:
            raise HTTPException(status_code=404, detail="task not found")
        if not subject.get("is_admin"):
            if str(row.get("user_id") or "") != str(subject["user_id"]):
                raise HTTPException(status_code=404, detail="task not found")
        attempts = db.list_attempts(task_id)
        latest_attempt = attempts[-1] if attempts else None
        desired = {
            "nnodes": int(row["nnodes"]),
            "n_gpus_per_node": int(row["n_gpus_per_node"]),
            "total_gpus": int(row["nnodes"]) * int(row["n_gpus_per_node"]),
        }
        out: dict[str, Any] = {
            "task_id": row["task_id"],
            "workload": row["workload"],
            "state": row["state"],
            "created_at": row.get("created_at"),
            "updated_at": row.get("updated_at"),
            "desired_resources": desired,
            "error_summary": row.get("error_summary"),
        }
        if latest_attempt:
            out["latest_attempt"] = {
                "attempt_no": latest_attempt["attempt_no"],
                "ray_submission_id": latest_attempt["ray_submission_id"],
                "ray_status": latest_attempt.get("ray_status"),
                "start_time": latest_attempt.get("start_time"),
                "end_time": latest_attempt.get("end_time"),
                "failure_kind": latest_attempt.get("failure_kind"),
                "message": latest_attempt.get("message"),
            }
        return out

    @app.get("/api/v2/tasks/{task_id}/attempts")
    async def get_attempts(task_id: str, req: Request) -> dict[str, Any]:
        subject = _auth(req)
        row = db.get_task(task_id)
        if not row:
            raise HTTPException(status_code=404, detail="task not found")
        if not subject.get("is_admin"):
            if str(row.get("user_id") or "") != str(subject["user_id"]):
                raise HTTPException(status_code=404, detail="task not found")
        return {"task_id": task_id, "attempts": db.list_attempts(task_id)}

    @app.post("/api/v2/tasks/{task_id}:cancel")
    async def cancel(task_id: str, req: Request) -> dict[str, Any]:
        subject = _auth(req)
        row = db.get_task(task_id)
        if not row:
            raise HTTPException(status_code=404, detail="task not found")
        if not subject.get("is_admin"):
            if str(row.get("user_id") or "") != str(subject["user_id"]):
                raise HTTPException(status_code=404, detail="task not found")

        state = str(row["state"])
        if state in ("SUCCEEDED", "FAILED", "CANCELED"):
            raise HTTPException(status_code=409, detail=f"task already terminal: {state}")

        attempts = db.list_attempts(task_id)
        if attempts:
            ray_sid = str(attempts[-1]["ray_submission_id"])
            try:
                tool.stop(ray_sid)
            except Exception:
                pass
            # Mark attempt as canceled on the service side so that API doesn't keep reporting RUNNING.
            # Ray stop is async; we deliberately reflect the user's intent here.
            db.update_attempt(
                task_id=task_id,
                attempt_no=int(attempts[-1]["attempt_no"]),
                ray_status="STOPPED",
                failure_kind="CANCELED",
                message="Canceled by user via API (Ray stop requested).",
                end_time=_utc_now_iso(),
            )

        db.set_task_state(task_id=task_id, state="CANCELED", event_type="CANCELED")
        return {"task_id": task_id, "state": "CANCELED"}

    @app.get("/api/v2/tasks/{task_id}/logs")
    async def logs(task_id: str, req: Request, tail: int = 2000, attempt: str = "latest") -> Response:
        subject = _auth(req)
        row = db.get_task(task_id)
        if not row:
            raise HTTPException(status_code=404, detail="task not found")
        if not subject.get("is_admin"):
            if str(row.get("user_id") or "") != str(subject["user_id"]):
                raise HTTPException(status_code=404, detail="task not found")
        attempts = db.list_attempts(task_id)
        if not attempts:
            raise HTTPException(status_code=404, detail="no attempts yet")
        a = attempts[-1] if attempt == "latest" else None
        if a is None:
            raise HTTPException(status_code=400, detail="only attempt=latest supported in v2.0")
        ray_sid = str(a["ray_submission_id"])
        text = tool.logs(ray_sid)
        if tail and tail > 0:
            lines = text.splitlines()
            text = "\n".join(lines[-tail:]) + ("\n" if lines else "")
        return Response(content=text, media_type="text/plain")

    @app.get("/api/v2/queue")
    async def queue(req: Request) -> dict[str, Any]:
        subject = _auth(req)
        if subject.get("is_admin"):
            return db.list_queue()
        return db.list_queue(user_id=str(subject["user_id"]))

    return app
