import base64
import tempfile
import time
import uuid
from http import HTTPStatus
from pathlib import Path
from typing import Any
from urllib.parse import unquote

from brainsurgery.core import IteratingTransform, IterationProgress, get_transform
from brainsurgery.engine import (
    RuntimeFlagLifecycleScope,
    list_model_aliases,
    reset_runtime_flags_for_scope,
)

from ..http import JsonRequestHandler
from .backend import (
    _api_error_payload,
    _apply_load_transform,
    _apply_transform,
    _default_alias,
    _model_diff_summary,
    _module_health_for_alias,
    _parse_filter_expr,
    _preview_transform,
    _render_dump_for_alias,
    _render_execution_summary,
    _require_string,
    _serialize_models,
    _serialize_runtime_flags,
    _transform_items,
)
from .models import (
    ApplyTransformResponsePayload,
    LoadResponsePayload,
    ModelDiffResponsePayload,
    ModelDumpResponsePayload,
    ProgressResponsePayload,
    SaveDownloadResponsePayload,
    StateResponsePayload,
    TransformsResponsePayload,
    model_to_payload,
)
from .page import _HTML_PAGE, _STATIC_DIR
from .session import _SessionState


def _now_ms() -> int:
    return int(time.time() * 1000)


def _snapshot_progress(session: _SessionState) -> dict[str, Any]:
    with session.progress_lock:
        return dict(session.progress)


def _begin_progress(session: _SessionState, *, transform: str, iterating: bool) -> None:
    now = _now_ms()
    with session.progress_lock:
        session.progress = {
            "active": True,
            "iterating": iterating,
            "transform": transform,
            "desc": None,
            "unit": "item",
            "completed": 0,
            "total": None,
            "started_at": now,
            "updated_at": now,
            "error": None,
        }


def _update_progress_from_event(session: _SessionState, event: IterationProgress) -> None:
    now = _now_ms()
    with session.progress_lock:
        session.progress["iterating"] = True
        session.progress["desc"] = event.desc
        session.progress["unit"] = event.unit
        session.progress["completed"] = int(event.completed)
        session.progress["total"] = int(event.total) if event.total is not None else None
        session.progress["updated_at"] = now


def _make_progress_callback(session: _SessionState):
    def _callback(event: IterationProgress) -> None:
        _update_progress_from_event(session, event)

    return _callback


def _finish_progress(session: _SessionState, *, error: str | None = None) -> None:
    now = _now_ms()
    with session.progress_lock:
        session.progress["active"] = False
        session.progress["updated_at"] = now
        session.progress["error"] = error


def _record_preview_history(session: _SessionState, entry: dict[str, Any]) -> None:
    entry_copy = dict(entry)
    entry_copy["timestamp_ms"] = _now_ms()
    session.preview_history.append(entry_copy)
    if len(session.preview_history) > 40:
        del session.preview_history[:-40]


def _serialize_insights(session: _SessionState) -> dict[str, Any]:
    return {"preview_history": list(session.preview_history)}


def _handler_factory(session: _SessionState):
    reset_runtime_flags_for_scope(RuntimeFlagLifecycleScope.WEBUI_SESSION)
    static_content_types = {
        ".html": "text/html; charset=utf-8",
        ".css": "text/css; charset=utf-8",
        ".js": "application/javascript; charset=utf-8",
    }

    class Handler(JsonRequestHandler):
        request_log_prefix = "webui"

        def _cache_headers(self) -> list[tuple[str, str]]:
            return [
                ("Cache-Control", "no-store, max-age=0"),
                ("Pragma", "no-cache"),
            ]

        def do_GET(self) -> None:  # noqa: N802
            if self.path == "/" or self.path.startswith("/?"):
                self._send_html(_HTML_PAGE)
                return
            if self.path.startswith("/static/"):
                static_rel = unquote(self.path[len("/static/") :].split("?", 1)[0]).lstrip("/")
                static_path = (_STATIC_DIR / static_rel).resolve()
                try:
                    static_path.relative_to(_STATIC_DIR.resolve())
                except ValueError:
                    self.send_error(HTTPStatus.NOT_FOUND, "Not Found")
                    return
                if not static_path.is_file():
                    self.send_error(HTTPStatus.NOT_FOUND, "Not Found")
                    return
                content_type = static_content_types.get(
                    static_path.suffix.lower(), "application/octet-stream"
                )
                encoded = static_path.read_bytes()
                self.send_response(HTTPStatus.OK)
                self.send_header("Content-Type", content_type)
                for key, value in self._cache_headers():
                    self.send_header(key, value)
                self.send_header("Content-Length", str(len(encoded)))
                self.end_headers()
                self.wfile.write(encoded)
                return
            if self.path == "/api/transforms":
                self._send_json(
                    model_to_payload(
                        TransformsResponsePayload(
                            ok=True,
                            transforms=_transform_items(),
                        )
                    )
                )
                return
            if self.path == "/api/state":
                with session.lock:
                    self._send_json(
                        model_to_payload(
                            StateResponsePayload(
                                ok=True,
                                models=_serialize_models(session.provider),
                                runtime_flags=_serialize_runtime_flags(),
                                insights=_serialize_insights(session),
                            )
                        )
                    )
                return
            if self.path == "/api/progress":
                self._send_json(
                    model_to_payload(
                        ProgressResponsePayload(
                            ok=True,
                            progress=_snapshot_progress(session),
                        )
                    )
                )
                return
            self.send_error(HTTPStatus.NOT_FOUND, "Not Found")

        def do_POST(self) -> None:  # noqa: N802
            if self.path == "/api/load":
                try:
                    body = self._read_json_body()
                    alias = body.get("alias")
                    if alias is not None and (not isinstance(alias, str) or not alias.strip()):
                        raise ValueError("alias must be a non-empty string when provided.")
                    alias_clean = alias.strip() if isinstance(alias, str) else None
                    server_path_raw = body.get("server_path")
                    server_path = (
                        server_path_raw.strip()
                        if isinstance(server_path_raw, str) and server_path_raw.strip()
                        else None
                    )

                    load_path: Path
                    if server_path is not None:
                        load_path = Path(server_path)
                    else:
                        filename = _require_string(body.get("filename"), "filename")
                        content_b64 = _require_string(body.get("content_b64"), "content_b64")
                        raw = base64.b64decode(content_b64, validate=True)
                        load_path = (
                            session.upload_root / f"{uuid.uuid4().hex}_{Path(filename).name}"
                        )
                        load_path.write_bytes(raw)
                except Exception as exc:
                    self._send_json({"ok": False, "error": str(exc)}, status=HTTPStatus.BAD_REQUEST)
                    return

                with session.lock:
                    try:
                        chosen_alias = alias_clean or _default_alias(session.provider)
                        _apply_load_transform(
                            provider=session.provider,
                            plan=session.plan,
                            path=load_path,
                            alias=chosen_alias,
                        )
                        session.default_model = chosen_alias
                        models = _serialize_models(session.provider)
                    except Exception as exc:
                        self._send_json(
                            {"ok": False, "error": str(exc)}, status=HTTPStatus.BAD_REQUEST
                        )
                        return

                self._send_json(
                    model_to_payload(
                        LoadResponsePayload(
                            ok=True,
                            models=models,
                            runtime_flags=_serialize_runtime_flags(),
                            insights=_serialize_insights(session),
                        )
                    )
                )
                return

            if self.path == "/api/_apply_transform":
                transform_name: str | None = None
                payload: Any | None = None
                preview_decision: str | None = None
                try:
                    body = self._read_json_body()
                    transform_name = _require_string(body.get("transform"), "transform")
                    payload = body.get("payload")
                    summary_mode_raw = body.get("summary_mode", "raw")
                    preview_decision_raw = body.get("preview_decision")
                    if payload is None:
                        payload = {}
                    if transform_name not in {"help", "assert"} and not isinstance(payload, dict):
                        raise ValueError("payload must be an object.")
                    if not isinstance(summary_mode_raw, str):
                        raise ValueError("summary_mode must be a string when provided.")
                    if preview_decision_raw is not None:
                        if not isinstance(preview_decision_raw, str):
                            raise ValueError("preview_decision must be a string when provided.")
                        normalized_preview_decision = preview_decision_raw.strip().lower()
                        if normalized_preview_decision not in {"go", "no-go"}:
                            raise ValueError("preview_decision must be 'go' or 'no-go'.")
                        preview_decision = normalized_preview_decision
                except Exception as exc:
                    self._send_json(
                        _api_error_payload(
                            exc,
                            endpoint="/api/_apply_transform",
                            transform_name=transform_name,
                            payload=payload,
                        ),
                        status=HTTPStatus.BAD_REQUEST,
                    )
                    return

                try:
                    iterating = isinstance(get_transform(transform_name), IteratingTransform)
                except Exception as exc:
                    self._send_json(
                        _api_error_payload(
                            exc,
                            endpoint="/api/_apply_transform",
                            transform_name=transform_name,
                            payload=payload,
                        ),
                        status=HTTPStatus.BAD_REQUEST,
                    )
                    return
                _begin_progress(session, transform=transform_name, iterating=iterating)
                progress_callback = _make_progress_callback(session) if iterating else None
                with session.lock:
                    try:
                        flags = _serialize_runtime_flags()
                        preview_enabled = bool(flags.preview)
                        dry_run_enabled = bool(flags.dry_run)
                        if (
                            transform_name != "set"
                            and preview_enabled
                            and not dry_run_enabled
                            and preview_decision != "go"
                        ):
                            (
                                preview_output,
                                needs_confirmation,
                                preview_impact,
                            ) = _preview_transform(
                                provider=session.provider,
                                transform_name=transform_name,
                                payload=payload,
                                default_model=session.default_model,
                            )
                            _record_preview_history(session, preview_impact)
                            models = _serialize_models(session.provider)
                            if needs_confirmation:
                                _finish_progress(session, error=None)
                                if preview_decision == "no-go":
                                    self._send_json(
                                        model_to_payload(
                                            ApplyTransformResponsePayload(
                                                ok=True,
                                                models=models,
                                                output=(
                                                    f"{preview_output}\n"
                                                    f"preview 1/1 {transform_name}: no-go, apply skipped"
                                                ),
                                                runtime_flags=_serialize_runtime_flags(),
                                                preview_confirmation_required=False,
                                                preview_transform=transform_name,
                                                insights=_serialize_insights(session),
                                            )
                                        )
                                    )
                                    return
                                self._send_json(
                                    model_to_payload(
                                        ApplyTransformResponsePayload(
                                            ok=True,
                                            models=models,
                                            output=preview_output,
                                            runtime_flags=_serialize_runtime_flags(),
                                            preview_confirmation_required=True,
                                            preview_transform=transform_name,
                                            insights=_serialize_insights(session),
                                        )
                                    )
                                )
                                return
                        output, next_default_model = _apply_transform(
                            provider=session.provider,
                            plan=session.plan,
                            transform_name=transform_name,
                            payload=payload,
                            default_model=session.default_model,
                            progress_callback=progress_callback,
                        )
                        session.default_model = next_default_model
                        if transform_name == "exit":
                            summary = _render_execution_summary(
                                plan=session.plan, mode=summary_mode_raw
                            )
                            output = f"{output}\n\n{summary}".strip() if output else summary
                        models = _serialize_models(session.provider)
                    except Exception as exc:
                        _finish_progress(session, error=str(exc))
                        self._send_json(
                            _api_error_payload(
                                exc,
                                endpoint="/api/_apply_transform",
                                transform_name=transform_name,
                                payload=payload,
                            ),
                            status=HTTPStatus.BAD_REQUEST,
                        )
                        return
                _finish_progress(session, error=None)
                self._send_json(
                    model_to_payload(
                        ApplyTransformResponsePayload(
                            ok=True,
                            models=models,
                            output=output,
                            runtime_flags=_serialize_runtime_flags(),
                            preview_confirmation_required=False,
                            preview_transform=None,
                            insights=_serialize_insights(session),
                        )
                    )
                )
                return

            if self.path == "/api/save_download":
                try:
                    body = self._read_json_body()
                    payload = body.get("payload")
                    if payload is None:
                        payload = {}
                    if not isinstance(payload, dict):
                        raise ValueError("payload must be an object.")
                except Exception as exc:
                    self._send_json({"ok": False, "error": str(exc)}, status=HTTPStatus.BAD_REQUEST)
                    return

                _begin_progress(session, transform="save", iterating=False)
                with session.lock:
                    try:
                        output, filename, mime, content_b64 = self._run_save_download(payload)
                        models = _serialize_models(session.provider)
                    except Exception as exc:
                        _finish_progress(session, error=str(exc))
                        self._send_json(
                            {"ok": False, "error": str(exc)}, status=HTTPStatus.BAD_REQUEST
                        )
                        return
                _finish_progress(session, error=None)
                self._send_json(
                    model_to_payload(
                        SaveDownloadResponsePayload(
                            ok=True,
                            models=models,
                            output=output,
                            runtime_flags=_serialize_runtime_flags(),
                            download_filename=filename,
                            download_mime=mime,
                            download_b64=content_b64,
                            insights=_serialize_insights(session),
                        )
                    )
                )
                return

            if self.path == "/api/model_dump":
                try:
                    body = self._read_json_body()
                    alias = _require_string(body.get("alias"), "alias")
                    format_name = _require_string(body.get("format"), "format").strip().lower()
                    if format_name not in {"compact", "tree"}:
                        raise ValueError("format must be 'compact' or 'tree'.")
                    verbosity = (
                        _require_string(body.get("verbosity", "shape"), "verbosity").strip().lower()
                    )
                    if verbosity not in {"shape", "stat"}:
                        raise ValueError("verbosity must be 'shape' or 'stat'.")
                    target = _parse_filter_expr(body.get("filter"), alias=alias)
                except Exception as exc:
                    self._send_json({"ok": False, "error": str(exc)}, status=HTTPStatus.BAD_REQUEST)
                    return

                with session.lock:
                    try:
                        if alias not in set(list_model_aliases(session.provider)):
                            raise ValueError(f"unknown alias: {alias!r}")
                        dumped, matched_count, total_count = _render_dump_for_alias(
                            provider=session.provider,
                            alias=alias,
                            format_name=format_name,
                            verbosity=verbosity,
                            target=target,
                        )
                        try:
                            module_health = _module_health_for_alias(
                                provider=session.provider,
                                alias=alias,
                                target=target,
                            )
                        except Exception:
                            module_health = None
                    except Exception as exc:
                        self._send_json(
                            {"ok": False, "error": str(exc)}, status=HTTPStatus.BAD_REQUEST
                        )
                        return

                self._send_json(
                    model_to_payload(
                        ModelDumpResponsePayload(
                            ok=True,
                            dump=dumped,
                            matched_count=matched_count,
                            total_count=total_count,
                            module_health=module_health,
                        )
                    )
                )
                return

            if self.path == "/api/model_diff":
                try:
                    body = self._read_json_body()
                    left_alias = _require_string(body.get("left_alias"), "left_alias")
                    right_alias = _require_string(body.get("right_alias"), "right_alias")
                    left_target = _parse_filter_expr(body.get("left_filter"), alias=left_alias)
                    right_target = _parse_filter_expr(body.get("right_filter"), alias=right_alias)
                    eps_raw = body.get("eps")
                    if eps_raw is None:
                        eps = None
                    elif isinstance(eps_raw, bool) or not isinstance(eps_raw, int | float):
                        raise ValueError("eps must be a non-negative number when provided.")
                    else:
                        eps = float(eps_raw)
                        if eps < 0:
                            raise ValueError("eps must be a non-negative number when provided.")
                except Exception as exc:
                    self._send_json({"ok": False, "error": str(exc)}, status=HTTPStatus.BAD_REQUEST)
                    return

                with session.lock:
                    try:
                        aliases = set(list_model_aliases(session.provider))
                        if left_alias not in aliases:
                            raise ValueError(f"unknown alias: {left_alias!r}")
                        if right_alias not in aliases:
                            raise ValueError(f"unknown alias: {right_alias!r}")
                        diff_payload = _model_diff_summary(
                            provider=session.provider,
                            left_alias=left_alias,
                            right_alias=right_alias,
                            left_target=left_target,
                            right_target=right_target,
                            eps=eps,
                        )
                    except Exception as exc:
                        self._send_json(
                            {"ok": False, "error": str(exc)}, status=HTTPStatus.BAD_REQUEST
                        )
                        return

                self._send_json(
                    model_to_payload(
                        ModelDiffResponsePayload(
                            ok=True,
                            diff=diff_payload,
                        )
                    )
                )
                return

            self.send_error(HTTPStatus.NOT_FOUND, "Not Found")

        def _run_save_download(self, payload: dict[str, Any]) -> tuple[str, str, str, str]:
            requested_path = _require_string(payload.get("path"), "payload.path")
            requested_name = Path(requested_path).name or "model"
            payload_copy = dict(payload)

            with tempfile.TemporaryDirectory(prefix="brainsurgery-webui-save-") as tmp:
                tmp_root = Path(tmp)
                out_path = tmp_root / requested_name
                payload_copy["path"] = str(out_path)

                output, next_default_model = _apply_transform(
                    provider=session.provider,
                    plan=session.plan,
                    transform_name="save",
                    payload=payload_copy,
                    default_model=session.default_model,
                    record_payload=payload,
                )
                session.default_model = next_default_model

                if out_path.is_dir():
                    raise ValueError(
                        "save download supports files only. "
                        "For directory outputs, use 'save on server path'."
                    )

                if out_path.is_file():
                    raw = out_path.read_bytes()
                    filename = _suggest_download_filename(
                        requested_name=requested_name,
                        out_path=out_path,
                        payload=payload_copy,
                    )
                    mime = "application/octet-stream"
                    return output, filename, mime, base64.b64encode(raw).decode("ascii")

                raise ValueError("save did not produce a file or directory to download.")

    return Handler


def _suggest_download_filename(
    *, requested_name: str, out_path: Path, payload: dict[str, Any]
) -> str:
    if out_path.suffix:
        return out_path.name

    base = requested_name or "model"
    if "." in base:
        return base

    fmt_raw = payload.get("format")
    fmt = fmt_raw.strip().lower() if isinstance(fmt_raw, str) else "safetensors"
    ext = {
        "safetensors": ".safetensors",
        "numpy": ".npy",
        "torch": ".pt",
    }.get(fmt, ".bin")
    return base + ext
