import json
import os
import platform
import secrets
import shutil
import socket
import subprocess
import sys
import tempfile
import threading
import time
import uuid
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.request import ProxyHandler, Request, build_opener

import bpy

from .models import (
    BEN2_MODEL,
    get_downloadable_model,
    model_directory_name,
    selected_model_key,
)


DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6"
RUNTIME_SCHEMA = 4
HOST_SCHEMA = 2
HOST_ADDRESS = "127.0.0.1"
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
JOB_FOLDER_NAMES = {
    "INSTALL_RUNTIME": "install-runtime",
    "DOWNLOAD_MODEL": "download-model",
    "RECONSTRUCT": "export-reconstruction",
    "IMAGE_DEPTH": "generate-depth-plane",
    "REMOVE_BACKGROUND": "remove-background",
    "IMAGE_BACKGROUND": "remove-image-background",
    "ORIGINAL_BACKGROUND": "generate-original-depth-plane",
    "ORIGINAL_FRAMES": "extract-original-frames",
}


@dataclass(frozen=True)
class HostJobRequest:
    payload: dict


@dataclass
class HostConnection:
    port: int
    token: str
    pid: int
    instance_id: str


@dataclass
class LaunchedHost:
    process: subprocess.Popen
    log_file: object


_host_connection = None
_launched_host = None
_registered_instance_id = None
_client_id = uuid.uuid4().hex
_host_auto_start = False
_host_snapshot = {"state": "STOPPED", "message": "Stopped"}
_local_opener = build_opener(ProxyHandler({}))
_client_label = f"Blender {getattr(bpy.app, 'version_string', '')}".strip()


def preferences():
    addon = bpy.context.preferences.addons.get(__package__)
    return addon.preferences if addon is not None else None


def storage_root():
    addon_preferences = preferences()
    configured = getattr(addon_preferences, "storage_root", "")
    root = (
        Path(bpy.path.abspath(configured)).expanduser()
        if configured
        else DEFAULT_STORAGE_ROOT
    )
    root.mkdir(parents=True, exist_ok=True)
    return root


def runtime_directory():
    return storage_root() / "runtime"


def runtime_python():
    directory = runtime_directory()
    if os.name == "nt":
        return directory / "Scripts" / "python.exe"
    return directory / "bin" / "python"


def runtime_ready():
    ready_path = runtime_directory() / "snapform_ready.json"
    if not runtime_python().is_file() or not ready_path.is_file():
        return False
    try:
        ready = json.loads(ready_path.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return False
    return ready.get("runtime_schema") == RUNTIME_SCHEMA


def models_directory():
    path = storage_root() / "models"
    path.mkdir(parents=True, exist_ok=True)
    return path


def model_directory(model_key):
    root = models_directory()
    directory = root / model_directory_name(model_key)
    legacy_directory = root / model_key.lower()
    if (
        not directory.is_dir()
        and legacy_directory.is_dir()
        and legacy_directory != directory
    ):
        legacy_directory.rename(directory)
    return directory


def model_ready(model_key):
    path = model_directory(model_key)
    return (path / "config.json").is_file() and any(
        path.glob("*.safetensors")
    )


def runtime_devices():
    """Devices detected by the installed Runtime, or None when unknown."""
    ready_path = runtime_directory() / "snapform_ready.json"
    if not ready_path.is_file():
        return None
    try:
        payload = json.loads(ready_path.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return None
    devices = payload.get("devices")
    if not isinstance(devices, list) or not devices:
        return None
    return [str(name).lower() for name in devices]


def machine_devices():
    """Devices the machine can run, fastest first; used before Runtime install."""
    devices = ["cpu"]
    if shutil.which("nvidia-smi"):
        devices.insert(0, "cuda")
    elif sys.platform == "darwin" and platform.machine().lower() in {
        "arm64",
        "aarch64",
    }:
        devices.insert(0, "mps")
    return devices


def supported_devices():
    """Inference devices available on this machine, fastest first."""
    return runtime_devices() or machine_devices()


DEVICE_LABELS = {
    "cuda": ("CUDA", "NVIDIA CUDA"),
    "mps": ("MPS", "Apple Metal (MPS)"),
    "cpu": ("CPU", "CPU"),
}


def device_enum_items(_owner, _context):
    """Blender EnumProperty items listing only devices supported on this machine."""
    return [
        (name.upper(), *DEVICE_LABELS[name], index)
        for index, name in enumerate(supported_devices())
    ]


def production_device():
    """Device configured for production tasks, falling back to Auto."""
    addon_preferences = preferences()
    configured = getattr(addon_preferences, "device", "")
    return configured or "AUTO"


def jobs_directory():
    path = storage_root() / "jobs"
    path.mkdir(parents=True, exist_ok=True)
    return path


def temporary_input_path():
    directory = Path(tempfile.mkdtemp(prefix="snapform-input-"))
    return directory / "input.png"


def temporary_input_directory():
    return Path(tempfile.mkdtemp(prefix="snapform-input-"))


def move_input_into_job(input_path, job_directory):
    """Move a Blender-generated input copy into the flat job directory."""
    input_path = Path(input_path)
    job_directory = Path(job_directory)
    if input_path.is_file():
        destination = job_directory / "input.png"
        input_path.replace(destination)
        try:
            input_path.parent.rmdir()
        except OSError:
            pass
        return destination
    if not input_path.is_dir():
        raise ValueError(f"Temporary input does not exist: {input_path}")
    for source in input_path.iterdir():
        if not source.is_file():
            raise ValueError(f"Temporary input contains a directory: {source}")
        destination = job_directory / source.name
        if destination.exists():
            raise ValueError(f"Job input filename already exists: {source.name}")
        source.replace(destination)
    input_path.rmdir()
    return job_directory


def job_folder_name(job_type):
    """Return the folder name segment for a job type."""
    return JOB_FOLDER_NAMES.get(job_type, job_type.lower())


def new_job_paths(job_type):
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
    stem = f"{job_folder_name(job_type)}_{timestamp}"
    directory = jobs_directory() / stem
    directory.mkdir(parents=True, exist_ok=True)
    return {
        "log": directory / f"{directory.name}.log",
        "status": directory / f"{directory.name}.json",
        "result": directory,
    }


def uv_executable():
    """Return the bundled uv path inside the shared storage."""
    return str(uv_cache_path())


def uv_cache_path():
    if os.name == "nt":
        return storage_root() / "uv" / "uv.exe"
    return storage_root() / "uv" / "uv"


def install_runtime_command(status_path):
    executable = uv_executable()
    backend = "cu128" if "cuda" in machine_devices() else "cpu"
    script = Path(__file__).with_name("bootstrap_runtime.py")
    command = [
        sys.executable,
        str(script),
        "--uv",
        executable,
        "--runtime-dir",
        str(runtime_directory()),
        "--models-dir",
        str(models_directory()),
        "--backend",
        backend,
        "--da3-commit",
        DA3_COMMIT,
        "--ben2-commit",
        BEN2_COMMIT,
        "--status",
        str(status_path),
    ]
    return command


def host_log_path():
    return storage_root() / "ai-host.log"


def host_discovery_path():
    return storage_root() / "host.json"


def host_lock_path():
    return storage_root() / "host.lock"


def _available_host_port():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
        listener.bind((HOST_ADDRESS, 0))
        return listener.getsockname()[1]


def _host_url(connection, path):
    return f"http://{HOST_ADDRESS}:{connection.port}{path}"


def _host_http_request(connection, method, path, payload=None, timeout=2.0):
    data = None
    headers = {"X-SnapForm-Token": connection.token}
    if payload is not None:
        data = json.dumps(payload).encode("utf-8")
        headers["Content-Type"] = "application/json"
    request = Request(
        _host_url(connection, path),
        data=data,
        headers=headers,
        method=method,
    )
    try:
        with _local_opener.open(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except HTTPError as error:
        try:
            detail = json.loads(error.read().decode("utf-8")).get("detail")
        except (json.JSONDecodeError, UnicodeDecodeError):
            detail = None
        raise RuntimeError(detail or f"AI Host returned HTTP {error.code}")
    except (OSError, URLError) as error:
        raise RuntimeError(f"Unable to reach AI Host: {error}")


def _read_host_discovery():
    try:
        payload = json.loads(host_discovery_path().read_text(encoding="utf-8"))
        return HostConnection(
            port=int(payload["port"]),
            token=str(payload["token"]),
            pid=int(payload["pid"]),
            instance_id=str(payload["instance_id"]),
        )
    except (KeyError, TypeError, ValueError, OSError, json.JSONDecodeError):
        return None


def _check_host(connection, timeout=0.3):
    health = _host_http_request(
        connection,
        "GET",
        "/health",
        timeout=timeout,
    )
    if health.get("host_schema") != HOST_SCHEMA:
        raise RuntimeError("AI Host protocol is out of date")
    if health.get("instance_id") != connection.instance_id:
        raise RuntimeError("AI Host discovery is stale")
    return health


def _discover_live_host(timeout=0.3):
    connection = _read_host_discovery()
    if connection is None:
        return None, None
    try:
        return connection, _check_host(connection, timeout=timeout)
    except RuntimeError:
        return None, None


def start_host():
    global _launched_host

    if _launched_host is not None and _launched_host.process.poll() is None:
        return
    _close_launched_host()
    port = _available_host_port()
    token = secrets.token_urlsafe(32)
    instance_id = uuid.uuid4().hex
    log_file = host_log_path().open("a", encoding="utf-8")
    command = [
        str(runtime_python()),
        "-u",
        str(Path(__file__).with_name("ai_host.py")),
        "--host",
        HOST_ADDRESS,
        "--port",
        str(port),
        "--token",
        token,
        "--instance-id",
        instance_id,
        "--discovery",
        str(host_discovery_path()),
        "--lock",
        str(host_lock_path()),
    ]
    try:
        process = subprocess.Popen(
            command,
            stdout=log_file,
            stderr=subprocess.STDOUT,
            cwd=storage_root(),
            creationflags=creation_flags(),
        )
    except OSError:
        log_file.close()
        raise
    _launched_host = LaunchedHost(process, log_file)


def ensure_host(timeout=60.0):
    global _host_auto_start, _host_connection, _host_snapshot

    _host_auto_start = True
    deadline = time.monotonic() + timeout
    last_launch = 0.0
    starting = False
    while time.monotonic() < deadline:
        connection, health = _discover_live_host(timeout=0.5)
        if connection is not None:
            _host_connection = connection
            _register_host_client(connection)
            _host_snapshot = _format_host_snapshot(connection, health)
            _redraw_host_ui()
            return connection
        if not starting:
            _host_snapshot = {"state": "STARTING", "message": "Starting"}
            _redraw_host_ui()
            starting = True
        _close_finished_launched_host()
        if (
            _launched_host is None
            and time.monotonic() - last_launch >= 1.0
        ):
            start_host()
            last_launch = time.monotonic()
        time.sleep(0.05)
    raise RuntimeError(f"AI Host did not become ready; log: {host_log_path()}")


def submit_host_job(request):
    connection = _host_connection or ensure_host()
    payload = dict(request.payload)
    payload["client_id"] = _client_id
    _host_http_request(
        connection,
        "POST",
        "/jobs",
        payload=payload,
        timeout=5.0,
    )
    return connection


def stop_host():
    global _host_auto_start, _host_connection, _host_snapshot
    global _registered_instance_id

    connection = _host_connection
    if connection is None:
        connection, _health = _discover_live_host(timeout=0.5)
    if connection is None:
        _close_launched_host(terminate=True)
        _host_auto_start = False
        _host_snapshot = {"state": "STOPPED", "message": "Stopped"}
        return
    _register_host_client(connection)
    _host_http_request(
        connection,
        "POST",
        "/shutdown",
        payload={"client_id": _client_id},
        timeout=2.0,
    )
    _host_auto_start = False
    deadline = time.monotonic() + 10.0
    while time.monotonic() < deadline:
        try:
            _check_host(connection, timeout=0.2)
        except RuntimeError:
            break
        time.sleep(0.05)
    else:
        raise RuntimeError("AI Host did not stop")
    _host_connection = None
    _registered_instance_id = None
    _close_finished_launched_host()
    _host_snapshot = {"state": "STOPPED", "message": "Stopped"}


def restart_host():
    global _host_auto_start, _host_connection, _registered_instance_id
    global _host_snapshot

    connection = ensure_host()
    _host_http_request(
        connection,
        "POST",
        "/shutdown",
        payload={"client_id": _client_id},
        timeout=2.0,
    )
    _host_auto_start = True
    _host_connection = None
    _registered_instance_id = None
    _host_snapshot = {"state": "STARTING", "message": "Restarting"}


def enable_host():
    global _host_auto_start, _host_connection, _host_snapshot

    _host_auto_start = True
    connection, health = _discover_live_host(timeout=0.2)
    if connection is not None:
        _host_connection = connection
        _register_host_client(connection)
        _host_snapshot = _format_host_snapshot(connection, health)
        _redraw_host_ui()
        return
    start_host()
    _host_snapshot = {"state": "STARTING", "message": "Starting"}
    _redraw_host_ui()


def clear_host_models():
    connection = ensure_host()
    return _host_http_request(
        connection,
        "POST",
        "/models/clear",
        timeout=5.0,
    )


def cancel_host_job(job_id):
    connection = ensure_host()
    return _host_http_request(
        connection,
        "DELETE",
        f"/jobs/{job_id}",
        payload={"client_id": _client_id},
        timeout=2.0,
    )


def detach_host_client():
    global _host_connection, _registered_instance_id

    connection = _host_connection
    if connection is not None and _registered_instance_id is not None:
        try:
            _host_http_request(
                connection,
                "DELETE",
                f"/clients/{_client_id}",
                timeout=0.5,
            )
        except RuntimeError:
            pass
    _host_connection = None
    _registered_instance_id = None
    # 如果 Host 是由当前 Blender 启动的，关闭 Blender 时直接终止它，无需等待空闲超时
    _close_launched_host(terminate=True)


def schedule_host_start():
    timers = getattr(bpy.app, "timers", None)
    if timers is None:
        return
    if timers.is_registered(_service_host_timer):
        return
    timers.register(_service_host_timer, first_interval=1.0, persistent=True)


def unschedule_host():
    timers = getattr(bpy.app, "timers", None)
    if timers is None:
        return
    if timers.is_registered(_service_host_timer):
        timers.unregister(_service_host_timer)


def host_status():
    return dict(_host_snapshot)


def host_busy():
    return _host_snapshot.get("state") == "BUSY"


def mark_host_job_complete(job_id):
    """Release a stale local busy snapshot after its terminal status is observed."""
    global _host_snapshot

    active = _host_snapshot.get("active_job")
    if active and active.get("job_id") != job_id:
        return
    if _host_snapshot.get("state") != "BUSY":
        return
    _host_snapshot = {
        **_host_snapshot,
        "busy": False,
        "active_job": None,
        "state": "READY",
        "message": "Ready",
    }


def _service_host_timer():
    global _host_connection, _host_snapshot, _registered_instance_id

    if not runtime_ready():
        _host_snapshot = {
            "state": "UNAVAILABLE",
            "message": "Runtime is not installed",
        }
        return 2.0
    connection, health = _discover_live_host(timeout=0.2)
    if connection is None:
        _host_connection = None
        _registered_instance_id = None
        _close_finished_launched_host()
        if not _host_auto_start:
            if _host_snapshot.get("state") != "STOPPED":
                _host_snapshot = {"state": "STOPPED", "message": "Stopped"}
                _redraw_host_ui()
            return 2.0
        try:
            start_host()
        except OSError as error:
            _host_snapshot = {"state": "ERROR", "message": str(error)}
            return 2.0
        _host_snapshot = {"state": "STARTING", "message": "Starting"}
        _redraw_host_ui()
        return 1.0
    try:
        _host_connection = connection
        if _registered_instance_id == connection.instance_id:
            try:
                _host_http_request(
                    connection,
                    "POST",
                    f"/clients/{_client_id}/heartbeat",
                    timeout=0.5,
                )
            except RuntimeError:
                _registered_instance_id = None
        _register_host_client(connection)
        health = _check_host(connection, timeout=0.5)
        _host_snapshot = _format_host_snapshot(connection, health)
    except RuntimeError as error:
        _host_connection = None
        _registered_instance_id = None
        _host_snapshot = {"state": "ERROR", "message": str(error)}
    _redraw_host_ui()
    return 5.0


def _register_host_client(connection):
    global _registered_instance_id

    if _registered_instance_id == connection.instance_id:
        return
    _host_http_request(
        connection,
        "POST",
        "/clients",
        payload={"client_id": _client_id, "label": _client_label},
        timeout=2.0,
    )
    _registered_instance_id = connection.instance_id


def _format_host_snapshot(connection, health):
    state = "BUSY" if health.get("busy") else "READY"
    return {
        **health,
        "state": state,
        "message": "Busy" if state == "BUSY" else "Ready",
        "port": connection.port,
    }


def _redraw_host_ui():
    if threading.current_thread() is not threading.main_thread():
        return
    try:
        from .status_bar import redraw_ui

        redraw_ui(bpy.context, force=True)
    except (AttributeError, RuntimeError):
        pass


def _close_finished_launched_host():
    if _launched_host is None or _launched_host.process.poll() is None:
        return
    _close_launched_host()


def _close_launched_host(terminate=False):
    global _launched_host

    launched = _launched_host
    _launched_host = None
    if launched is None:
        return
    if terminate and launched.process.poll() is None:
        launched.process.terminate()
        try:
            launched.process.wait(timeout=3.0)
        except subprocess.TimeoutExpired:
            launched.process.kill()
            launched.process.wait(timeout=3.0)
    launched.log_file.close()


def download_model_command(model_key, status_path):
    model = get_downloadable_model(model_key)
    command = [
        str(runtime_python()),
        str(Path(__file__).with_name("worker.py")),
        "download",
        "--repository",
        model.repository,
        "--model-dir",
        str(model_directory(model_key)),
        "--status",
        str(status_path),
        "--label",
        model.label,
    ]
    for pattern in model.download_patterns:
        command.extend(["--allow-pattern", pattern])
    return command


def remove_background_request(
    settings,
    input_path,
    status_path,
    result_directory,
    model_key,
    video_frames=None,
    start_frame=1,
    device=None,
    background_depth_distance=0.0,
):
    return HostJobRequest(
        {
            "job_id": Path(status_path).stem,
            "command": "remove-background",
            "parameters": {
                "input": str(input_path),
                "model_dir": str(model_directory(model_key)),
                "output_dir": str(result_directory),
                "status": str(status_path),
                "device": (device or settings.device).lower(),
                "video_frames": (
                    settings.video_frames
                    if video_frames is None
                    else video_frames
                ),
                "start_frame": start_frame,
                "background_depth_distance": float(
                    background_depth_distance
                ),
            },
        }
    )


def extract_frames_request(
    settings,
    input_path,
    status_path,
    result_directory,
    video_frames=None,
    start_frame=1,
):
    return HostJobRequest(
        {
            "job_id": Path(status_path).stem,
            "command": "extract-frames",
            "parameters": {
                "input": str(input_path),
                "output_dir": str(result_directory),
                "status": str(status_path),
                "video_frames": (
                    settings.video_frames
                    if video_frames is None
                    else video_frames
                ),
                "start_frame": start_frame,
            },
        }
    )


def reconstruct_request(
    settings,
    export_format,
    status_path,
    result_directory,
    input_path=None,
    model_key=None,
    process_res=None,
    process_res_method=None,
    use_ray_pose=None,
    remove_background=False,
    background_depth_distance=0.0,
    start_frame=1,
    device=None,
):
    snapform_selected_key = model_key or selected_model_key(settings, export_format)
    snapform_selected_process_res = (
        settings.process_res if process_res is None else process_res
    )
    snapform_selected_process_res_method = (
        settings.process_res_method
        if process_res_method is None
        else process_res_method
    )
    snapform_selected_ray_pose = (
        settings.use_ray_pose if use_ray_pose is None else use_ray_pose
    )
    parameters = {
        "input": str(
            Path(input_path)
            if input_path is not None
            else Path(bpy.path.abspath(settings.input_path))
        ),
        "model_dir": str(model_directory(snapform_selected_key)),
        "output_dir": str(result_directory),
        "export_format": export_format,
        "status": str(status_path),
        "process_res": snapform_selected_process_res,
        "process_res_method": snapform_selected_process_res_method,
        "conf_thresh_percentile": settings.conf_thresh_percentile,
        "num_max_points": settings.num_max_points,
        "show_cameras": 1 if settings.show_cameras else 0,
        "gs_views_interval": settings.gs_views_interval,
        "video_frames": settings.video_frames,
        "start_frame": start_frame,
        "device": (device or settings.device).lower(),
        "use_ray_pose": 1 if snapform_selected_ray_pose else 0,
        "ref_view_strategy": settings.ref_view_strategy,
        "remove_background_model_dir": "",
        "background_depth_distance": float(background_depth_distance),
        "mesh_filter_edges": bool(
            getattr(settings, "mesh_filter_edges", True)
        ),
        "mesh_max_edge_length": float(
            getattr(settings, "mesh_max_edge_length", 0.1)
        ),
    }
    if remove_background or background_depth_distance > 0.0:
        parameters["remove_background_model_dir"] = str(
            model_directory(BEN2_MODEL.key)
        )
    return HostJobRequest(
        {
            "job_id": Path(status_path).stem,
            "command": "reconstruct",
            "parameters": parameters,
        }
    )


def creation_flags():
    if os.name != "nt":
        return 0
    return subprocess.CREATE_NO_WINDOW


def read_status(path):
    try:
        return json.loads(Path(path).read_text(encoding="utf-8"))
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        return None