from __future__ import annotations

import subprocess
import shutil
import uuid
from pathlib import Path
from typing import Any

import cv2
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, Response
from fastapi.staticfiles import StaticFiles

from backend.segmentation import METHOD_DESCRIPTIONS, compare_frame, segment_frame


ROOT = Path(__file__).resolve().parents[1]
FRONTEND_DIR = ROOT / "frontend"
STORAGE_DIR = ROOT / "storage"
UPLOAD_DIR = STORAGE_DIR / "uploads"
JOB_DIR = STORAGE_DIR / "jobs"
SAMPLE_DIR = STORAGE_DIR / "samples"

IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}
VIDEO_SUFFIXES = {".mp4", ".avi", ".mov", ".mkv", ".webm"}

app = FastAPI(title="ISISeg Guidewire Segmentation", version="0.1.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


def ensure_dirs() -> None:
    for directory in (UPLOAD_DIR, JOB_DIR, SAMPLE_DIR):
        directory.mkdir(parents=True, exist_ok=True)


ensure_dirs()
app.mount("/storage", StaticFiles(directory=STORAGE_DIR), name="storage")


@app.get("/api/health")
def health() -> dict[str, str]:
    return {"status": "ok", "service": "ISISeg", "version": app.version}


@app.get("/api/methods")
def methods() -> dict[str, Any]:
    return {"methods": METHOD_DESCRIPTIONS}


@app.get("/api/samples")
def samples() -> dict[str, Any]:
    ensure_dirs()
    items = []
    for path in sorted(SAMPLE_DIR.glob("*")):
        suffix = path.suffix.lower()
        if suffix in IMAGE_SUFFIXES | VIDEO_SUFFIXES:
            version = str(path.stat().st_mtime_ns)
            items.append(
                {
                    "name": path.name,
                    "url": f"{_public(path)}?v={version}",
                    "kind": "image" if suffix in IMAGE_SUFFIXES else "video",
                    "size": path.stat().st_size,
                    "version": version,
                }
            )
    return {"samples": items}


@app.get("/favicon.ico", include_in_schema=False)
def favicon() -> Response:
    svg = """<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 64 64">
<rect width="64" height="64" rx="12" fill="#10211d"/>
<path d="M18 40V18h6v22h-6Zm14 0c-4.9 0-8.2-2.4-8.5-6.4h5.8c.3 1.4 1.4 2.2 3.1 2.2 1.6 0 2.6-.6 2.6-1.8 0-1.3-1.3-1.6-4.3-2.3-3.4-.8-6.6-2-6.6-6.5 0-4.3 3.3-7.2 8.1-7.2 4.7 0 7.8 2.4 8.2 6.4h-5.7c-.2-1.2-1.1-1.9-2.6-1.9-1.4 0-2.3.6-2.3 1.7 0 1.1 1.2 1.5 3.8 2.1 3.8.9 7.2 2.1 7.2 6.9 0 4.2-3.4 6.8-8.8 6.8Z" fill="#38d8b8"/>
</svg>"""
    return Response(content=svg, media_type="image/svg+xml")


def _public(path: Path) -> str:
    return "/" + path.relative_to(ROOT).as_posix()


def _save_upload(file: UploadFile, job_path: Path) -> Path:
    suffix = Path(file.filename or "").suffix.lower()
    if suffix not in IMAGE_SUFFIXES | VIDEO_SUFFIXES:
        raise HTTPException(status_code=400, detail="仅支持图片或视频文件")
    destination = job_path / f"upload{suffix}"
    with destination.open("wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    return destination


def _save_frame_outputs(
    job_path: Path,
    frame_index: int,
    method: str,
    frame,
    output,
    source_time: float | None = None,
    result_time: float | None = None,
    result_index: int | None = None,
) -> dict[str, Any]:
    method_path = job_path / method
    method_path.mkdir(exist_ok=True)
    original_path = method_path / f"frame_{frame_index:04d}_original.png"
    mask_path = method_path / f"frame_{frame_index:04d}_mask.png"
    overlay_path = method_path / f"frame_{frame_index:04d}_overlay.png"
    cv2.imwrite(str(original_path), frame)
    cv2.imwrite(str(mask_path), output.mask)
    cv2.imwrite(str(overlay_path), output.overlay)
    payload = {
        "frame_index": frame_index,
        "method": method,
        "original_url": _public(original_path),
        "mask_url": _public(mask_path),
        "overlay_url": _public(overlay_path),
        "metrics": output.metrics,
    }
    if source_time is not None:
        payload["source_time"] = round(float(source_time), 4)
    if result_time is not None:
        payload["result_time"] = round(float(result_time), 4)
    if result_index is not None:
        payload["result_index"] = int(result_index)
    return payload


def _browser_video(raw_path: Path, final_path: Path) -> Path:
    ffmpeg = shutil.which("ffmpeg")
    if ffmpeg:
        subprocess.run(
            [
                ffmpeg,
                "-y",
                "-i",
                str(raw_path),
                "-c:v",
                "libx264",
                "-pix_fmt",
                "yuv420p",
                "-movflags",
                "+faststart",
                "-preset",
                "veryfast",
                "-crf",
                "20",
                str(final_path),
            ],
            check=True,
            capture_output=True,
            text=True,
        )
        raw_path.unlink(missing_ok=True)
    else:
        raw_path.replace(final_path)
    return final_path


def _process_image(job_path: Path, source: Path, method: str, sensitivity: float) -> dict[str, Any]:
    frame = cv2.imread(str(source), cv2.IMREAD_COLOR)
    if frame is None:
        raise HTTPException(status_code=400, detail="无法读取图片")
    if method == "compare":
        results = [
            _save_frame_outputs(job_path, 0, item.method, frame, item, 0.0, 0.0, 0)
            for item in compare_frame(frame, None, sensitivity)
        ]
    else:
        output = segment_frame(frame, method, None, sensitivity)
        results = [_save_frame_outputs(job_path, 0, method, frame, output, 0.0, 0.0, 0)]
    return {"kind": "image", "frames": results, "video_url": None, "source_fps": 1.0, "result_fps": 1.0, "duration": 0.0}


def _selected_frame(index: int, stride: int, selected_count: int, max_frames: int) -> bool:
    return index % max(1, stride) == 0 and selected_count < max_frames


def _selected_frame_indices(frame_count: int, stride: int, max_frames: int) -> set[int]:
    if frame_count <= 0:
        return set()
    candidates = list(range(0, frame_count, max(1, stride)))
    if len(candidates) <= max_frames:
        return set(candidates)
    if max_frames <= 1:
        return {candidates[0]}
    last = len(candidates) - 1
    return {
        candidates[round(position * last / (max_frames - 1))]
        for position in range(max_frames)
    }


def _process_video(
    job_path: Path,
    source: Path,
    method: str,
    sensitivity: float,
    frame_stride: int,
    max_frames: int,
) -> dict[str, Any]:
    capture = cv2.VideoCapture(str(source))
    if not capture.isOpened():
        raise HTTPException(status_code=400, detail="无法读取视频")

    frames: list[dict[str, Any]] = []
    previous = None
    frame_index = 0
    selected_count = 0
    written_count = 0
    writer = None
    raw_video_path = job_path / f"{method}_overlay.raw.mp4"
    video_path = job_path / f"{method}_overlay.mp4"
    source_fps = float(capture.get(cv2.CAP_PROP_FPS) or 0.0)
    if source_fps <= 0:
        source_fps = 12.0
    result_fps = source_fps
    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    selected_indices = _selected_frame_indices(frame_count, frame_stride, max_frames)

    try:
        while True:
            ok, frame = capture.read()
            if not ok:
                break
            should_process = (
                frame_index in selected_indices
                if selected_indices
                else _selected_frame(frame_index, frame_stride, selected_count, max_frames)
            )
            video_frame = frame

            if should_process and method == "compare":
                outputs = compare_frame(frame, previous, sensitivity)
                for output in outputs:
                    frames.append(
                        _save_frame_outputs(
                            job_path,
                            frame_index,
                            output.method,
                            frame,
                            output,
                            frame_index / source_fps,
                            frame_index / source_fps,
                            frame_index,
                        )
                    )
                video_output = next(item for item in outputs if item.method == "fusion")
                video_frame = video_output.overlay
                selected_count += 1
            elif should_process:
                video_output = segment_frame(frame, method, previous, sensitivity)
                frames.append(
                    _save_frame_outputs(
                        job_path,
                        frame_index,
                        method,
                        frame,
                        video_output,
                        frame_index / source_fps,
                        frame_index / source_fps,
                        frame_index,
                    )
                )
                video_frame = video_output.overlay
                selected_count += 1

            if writer is None:
                height, width = frame.shape[:2]
                fourcc = cv2.VideoWriter_fourcc(*"mp4v")
                writer = cv2.VideoWriter(str(raw_video_path), fourcc, result_fps, (width, height))
            writer.write(video_frame)
            written_count += 1

            previous = frame
            frame_index += 1
    finally:
        capture.release()
        if writer is not None:
            writer.release()

    if not frames:
        raise HTTPException(status_code=400, detail="视频没有可处理帧")
    if raw_video_path.exists():
        _browser_video(raw_video_path, video_path)
    duration = round(written_count / source_fps, 4) if written_count else 0.0
    return {
        "kind": "video",
        "frames": frames,
        "video_url": _public(video_path) if video_path.exists() else None,
        "source_fps": round(source_fps, 4),
        "result_fps": round(result_fps, 4),
        "duration": duration,
        "result_duration": duration,
    }


def _read_video_frame(source: Path, frame_index: int) -> tuple[Any, Any | None, float]:
    capture = cv2.VideoCapture(str(source))
    if not capture.isOpened():
        raise HTTPException(status_code=400, detail="无法读取视频")
    source_fps = float(capture.get(cv2.CAP_PROP_FPS) or 0.0)
    if source_fps <= 0:
        source_fps = 12.0
    frame_index = max(0, int(frame_index))
    previous = None
    if frame_index > 0:
        capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index - 1)
        ok_prev, previous = capture.read()
        if not ok_prev:
            previous = None
    capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
    ok, frame = capture.read()
    capture.release()
    if not ok:
        raise HTTPException(status_code=400, detail="无法读取指定帧")
    return frame, previous, source_fps


@app.post("/api/segment")
def segment(
    file: UploadFile = File(...),
    method: str = Form("fusion"),
    sensitivity: float = Form(0.56),
    frame_stride: int = Form(5),
    max_frames: int = Form(12),
) -> dict[str, Any]:
    ensure_dirs()
    if method not in METHOD_DESCRIPTIONS:
        raise HTTPException(status_code=400, detail="未知分割方法")
    sensitivity = max(0.05, min(float(sensitivity), 0.95))
    frame_stride = max(1, min(int(frame_stride), 90))
    max_frames = max(1, min(int(max_frames), 80))
    job_id = uuid.uuid4().hex[:12]
    job_path = JOB_DIR / job_id
    job_path.mkdir(parents=True, exist_ok=True)
    source = _save_upload(file, job_path)
    suffix = source.suffix.lower()
    if suffix in IMAGE_SUFFIXES:
        result = _process_image(job_path, source, method, sensitivity)
    elif suffix in VIDEO_SUFFIXES:
        result = _process_video(job_path, source, method, sensitivity, frame_stride, max_frames)
    else:
        raise HTTPException(status_code=400, detail="不支持的文件类型")
    return {
        "job_id": job_id,
        "method": method,
        "sensitivity": sensitivity,
        "frame_stride": frame_stride,
        "max_frames": max_frames,
        **result,
    }


@app.post("/api/compare-frame")
def compare_single_frame(
    file: UploadFile = File(...),
    frame_index: int = Form(0),
    sensitivity: float = Form(0.56),
) -> dict[str, Any]:
    ensure_dirs()
    sensitivity = max(0.05, min(float(sensitivity), 0.95))
    job_id = uuid.uuid4().hex[:12]
    job_path = JOB_DIR / job_id
    job_path.mkdir(parents=True, exist_ok=True)
    source = _save_upload(file, job_path)
    suffix = source.suffix.lower()
    if suffix in IMAGE_SUFFIXES:
        frame = cv2.imread(str(source), cv2.IMREAD_COLOR)
        previous = None
        source_time = 0.0
    elif suffix in VIDEO_SUFFIXES:
        frame, previous, source_fps = _read_video_frame(source, frame_index)
        source_time = int(frame_index) / source_fps
    else:
        raise HTTPException(status_code=400, detail="不支持的文件类型")
    if frame is None:
        raise HTTPException(status_code=400, detail="无法读取文件")
    outputs = compare_frame(frame, previous, sensitivity)
    frames = [
        _save_frame_outputs(job_path, int(frame_index), output.method, frame, output, source_time, 0.0, 0)
        for output in outputs
    ]
    return {
        "job_id": job_id,
        "kind": "compare",
        "frame_index": int(frame_index),
        "source_time": round(float(source_time), 4),
        "frames": frames,
    }


@app.get("/")
def index() -> FileResponse:
    return FileResponse(FRONTEND_DIR / "index.html")


app.mount("/", StaticFiles(directory=FRONTEND_DIR), name="frontend")