from __future__ import annotations

import argparse
import asyncio
import json
import os
import signal
import socket
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Any, AsyncGenerator, Dict, List

import requests
import uvicorn
from fastapi import FastAPI, HTTPException, Request, Response, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse

from daisy_sdk.constants import (
    CONTROL_INTERNAL_HOST,
    DEFAULT_APP_PORT,
    DEFAULT_CONTAINER_NAME,
    DEFAULT_CONTROL_BIND,
    DEFAULT_CONTROL_HOST,
    DEFAULT_CONTROL_PORT,
)
from daisy_sdk.container import (
    DockerError,
    container_running,
    describe_container,
    project_mounts,
    start_container,
    stop_container,
)
from daisy_sdk.events import EventBroadcaster
from daisy_sdk.utils import DEFAULT_PROJECT_CONFIG_PATH, read_toml, write_toml


class ControlContext:
    def __init__(
        self,
        image: str,
        app_port: str = DEFAULT_APP_PORT,
        project_config_path: Path = DEFAULT_PROJECT_CONFIG_PATH,
        platform: str | None = None,
        container_name: str = DEFAULT_CONTAINER_NAME,
        control_host: str = DEFAULT_CONTROL_HOST,
        control_port: int = DEFAULT_CONTROL_PORT,
        log_container: bool = False,
        show_docker_command: bool = False,
    ) -> None:
        self.image = image
        self.app_port = app_port
        self.project_config_path = project_config_path
        self.platform = platform
        self.container_name = container_name
        self.control_host = control_host
        self.control_port = control_port
        self.broadcaster = EventBroadcaster()
        self._lock = threading.Lock()
        self.log_process: subprocess.Popen[str] | None = None
        self.log_container = log_container
        self.show_docker_command = show_docker_command

    def update_from_payload(self, payload: Dict[str, Any]) -> None:
        with self._lock:
            image = payload.get("image")
            if isinstance(image, str) and image:
                self.image = image
            port = payload.get("port")
            if isinstance(port, str) and port:
                self.app_port = port
            platform = payload.get("platform")
            if isinstance(platform, str) and platform:
                self.platform = platform


app: FastAPI | None = None
context: ControlContext | None = None


def _ensure_context() -> ControlContext:
    if context is None:
        raise RuntimeError("Control context not initialized")
    return context


def _start_log_forwarder(ctx: ControlContext) -> None:
    try:
        proc = subprocess.Popen(
            ["docker", "logs", "-f", ctx.container_name],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        )
    except Exception:
        return

    def _pump() -> None:
        if not proc.stdout:
            return
        for line in proc.stdout:
            print(f"[daisy-container] {line.rstrip()}", flush=True)
        proc.wait()

    thread = threading.Thread(target=_pump, daemon=True)
    thread.start()
    ctx.log_process = proc


def _stop_log_forwarder(ctx: ControlContext) -> None:
    proc = ctx.log_process
    if proc:
        try:
            proc.terminate()
        except Exception:
            pass
    ctx.log_process = None


def _validate_path(path: str) -> tuple[bool, str]:
    if not isinstance(path, str) or not path:
        return False, "path is required"
    p = Path(path)
    if not p.is_absolute():
        return False, "path must be absolute"
    if not p.exists():
        return False, f"path does not exist: {path}"
    return True, ""


def _load_projects(cfg_path: Path) -> Dict[str, Any]:
    data = read_toml(cfg_path)
    # For legacy format we only care about top-level sections (no projects list).
    if isinstance(data.get("projects"), list):
        data.pop("projects", None)
    return data


def _persist_projects(cfg_path: Path, data: Dict[str, Any]) -> None:
    data = dict(data)
    # Never write a projects list; keep legacy top-level sections only.
    data.pop("projects", None)
    write_toml(cfg_path, data)


def _update_project_entry(entry: Dict[str, Any], payload: Dict[str, Any]) -> Dict[str, Any]:
    updated = dict(entry)
    # Map incoming fields to legacy keys.
    if "path" in payload or "project_path" in payload:
        path = payload.get("project_path") or payload.get("path")
        if isinstance(path, str):
            updated["project_path"] = path
    if "bi_paths" in payload and isinstance(payload["bi_paths"], list):
        updated["bi_paths"] = payload["bi_paths"]
    if "project_type" in payload and isinstance(payload["project_type"], str):
        updated["project_type"] = payload["project_type"]
    if "profile_target" in payload and isinstance(payload["profile_target"], str):
        updated["profile_target"] = payload["profile_target"]
    return updated


def _wait_for_container_ready(ctx: ControlContext, timeout: float = 20.0, interval: float = 0.5) -> tuple[bool, str]:
    """Simple fixed delay; skip polling health."""
    time.sleep(1.7)
    return True, "ready"


def _start_and_report(ctx: ControlContext) -> tuple[bool, str]:
    try:
        stop_container(ctx.container_name)
        _stop_log_forwarder(ctx)
        start_container(
            ctx.image,
            port=ctx.app_port,
            project_config_path=ctx.project_config_path,
            platform=ctx.platform,
            container_name=ctx.container_name,
            env={
                "DAISY_CONTROL_HOST": ctx.control_host,
                "DAISY_CONTROL_PORT": str(ctx.control_port),
            },
            pull=False,
            print_command=ctx.show_docker_command,
        )
        if ctx.log_container:
            _start_log_forwarder(ctx)
        ok, msg = _wait_for_container_ready(ctx)
        if not ok:
            return False, msg
        return True, "restarted"
    except DockerError as exc:
        return False, str(exc)


def _setup_app() -> FastAPI:
    api = FastAPI()
    api.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
    )

    @api.get("/health")
    async def health() -> Dict[str, Any]:
        ctx = _ensure_context()
        return {
            "status": "ok",
            "running": container_running(ctx.container_name),
            "container": describe_container(ctx.container_name),
            "image": ctx.image,
            "container_name": ctx.container_name,
            "mounts": project_mounts(ctx.project_config_path),
        }

    @api.post("/restart")
    async def restart(payload: Dict[str, Any] | None = None) -> JSONResponse:
        ctx = _ensure_context()
        payload = payload or {}
        ctx.update_from_payload(payload)
        success, message = _start_and_report(ctx)
        if success:
            event: Dict[str, Any] = {"type": "refresh"}
            project_path = payload.get("project_path")
            if isinstance(project_path, str) and project_path:
                event["activeProject"] = project_path
            ctx.broadcaster.publish(event)
            return JSONResponse({"status": "ok", "message": message})
        raise HTTPException(status_code=500, detail=message)

    @api.post("/projects")
    async def create_project(payload: Dict[str, Any]) -> JSONResponse:
        ctx = _ensure_context()
        name = payload.get("name")
        path = payload.get("project_path") or payload.get("path")
        if not isinstance(name, str) or not name:
            raise HTTPException(status_code=400, detail="name is required")
        ok, msg = _validate_path(path)
        if not ok:
            raise HTTPException(status_code=400, detail=msg)
        data = _load_projects(ctx.project_config_path)
        entry: Dict[str, Any] = {
            "project_path": path,
            "project_type": payload.get("project_type"),
            "profile_target": payload.get("profile_target"),
        }
        if isinstance(payload.get("bi_paths"), list):
            entry["bi_paths"] = payload["bi_paths"]

        legacy_section = data.get(name)
        if isinstance(legacy_section, dict):
            legacy_section.update({k: v for k, v in entry.items() if v is not None})
            data[name] = legacy_section
        else:
            data[name] = {k: v for k, v in entry.items() if v is not None}
        _persist_projects(ctx.project_config_path, data)
        success, message = _start_and_report(ctx)
        ctx.broadcaster.publish({"type": "refresh", "activeProject": path})
        if not success:
            raise HTTPException(status_code=500, detail=message)
        return JSONResponse({"status": "ok", "message": message, "project": entry})

    @api.patch("/projects/{project_name}")
    async def update_project(project_name: str, payload: Dict[str, Any]) -> JSONResponse:
        ctx = _ensure_context()
        data = _load_projects(ctx.project_config_path)
        legacy_section = data.get(project_name)

        if isinstance(legacy_section, dict):
            candidate = dict(legacy_section)
            if "path" in payload or "project_path" in payload:
                new_path = payload.get("project_path") or payload.get("path")
                ok, msg = _validate_path(new_path)
                if not ok:
                    raise HTTPException(status_code=400, detail=msg)
                candidate["project_path"] = new_path
            if "bi_paths" in payload and isinstance(payload["bi_paths"], list):
                candidate["bi_paths"] = payload["bi_paths"]
            if "project_type" in payload and isinstance(payload["project_type"], str):
                candidate["project_type"] = payload["project_type"]
            if "profile_target" in payload and isinstance(payload["profile_target"], str):
                candidate["profile_target"] = payload["profile_target"]
            data[project_name] = candidate
            active_path = candidate.get("project_path")
            persisted_project = {"name": project_name, "path": candidate.get("project_path")}
        else:
            raise HTTPException(status_code=404, detail=f"project '{project_name}' not found")

        _persist_projects(ctx.project_config_path, data)
        success, message = _start_and_report(ctx)
        event: Dict[str, Any] = {"type": "refresh"}
        if isinstance(active_path, str):
            event["activeProject"] = active_path
        ctx.broadcaster.publish(event)
        if not success:
            raise HTTPException(status_code=500, detail=message)
        return JSONResponse({"status": "ok", "message": message, "project": persisted_project})

    @api.get("/events")
    async def events() -> StreamingResponse:
        ctx = _ensure_context()
        queue = ctx.broadcaster.subscribe()

        async def event_generator() -> AsyncGenerator[str, None]:
            try:
                while True:
                    try:
                        data = queue.get_nowait()
                    except Exception:
                        await asyncio.sleep(0.5)
                        continue
                    yield f"data: {data}\n\n"
            finally:
                ctx.broadcaster.unsubscribe(queue)

        return StreamingResponse(event_generator(), media_type="text/event-stream")

    @api.post("/shutdown")
    async def shutdown_endpoint() -> JSONResponse:
        ctx = _ensure_context()
        stop_container(ctx.container_name)
        _stop_log_forwarder(ctx)
        def _stop_loop() -> None:
            loop = asyncio.get_event_loop()
            loop.stop()
        threading.Thread(target=_stop_loop, daemon=True).start()
        return JSONResponse({"status": "ok", "message": "shutting down"})

    return api


def run_server(
    *,
    bind_host: str = DEFAULT_CONTROL_BIND,
    client_host: str = DEFAULT_CONTROL_HOST,
    port: int = DEFAULT_CONTROL_PORT,
    image: str,
    app_port: str = DEFAULT_APP_PORT,
    project_config_path: Path = DEFAULT_PROJECT_CONFIG_PATH,
    platform: str | None = None,
    container_name: str = DEFAULT_CONTAINER_NAME,
    detach: bool = False,
    control_host_env: str = CONTROL_INTERNAL_HOST,
    log_container: bool = False,
    show_docker_command: bool = False,
) -> None:
    # Port check for default port
    if port == DEFAULT_CONTROL_PORT:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            try:
                sock.bind((bind_host, port))
            except OSError:
                print(
                    f"Port {DEFAULT_CONTROL_PORT} in use; Daisy image expects {DEFAULT_CONTROL_PORT}. "
                    "Free the port or rebuild the image with a new VITE_HOST_API_BASE."
                )
                return

    global app, context
    context = ControlContext(
        image=image,
        app_port=app_port,
        project_config_path=project_config_path,
        platform=platform,
        container_name=container_name,
        control_host=control_host_env,
        control_port=port,
        log_container=log_container,
        show_docker_command=show_docker_command,
    )
    app = _setup_app()

    # Initial start to ensure container is running before serving API
    started, message = _start_and_report(context)
    if not started:
        print(f"[daisy-control] Failed to start container: {message}")
        sys.exit(1)

    if detach:
        cmd = [
            sys.executable,
            "-m",
            "daisy_sdk.server",
            "--host",
            bind_host,
            "--port",
            str(port),
            "--image",
            image,
            "--app-port",
            app_port,
            "--project-config-path",
            str(project_config_path),
            *(["--platform", platform] if platform else []),
            "--container-name",
            container_name,
            "--container-logs" if log_container else "",
            "--show-docker-command" if show_docker_command else "",
        ]
        subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        return

    config = uvicorn.Config(app, host=bind_host, port=port, log_level="info")
    server = uvicorn.Server(config)

    def _handle_sig(*_: Any) -> None:
        try:
            stop_container(container_name)
            _stop_log_forwarder(context)
        finally:
            server.should_exit = True

    signal.signal(signal.SIGTERM, _handle_sig)
    signal.signal(signal.SIGINT, _handle_sig)

    server.run()


def main() -> None:
    parser = argparse.ArgumentParser(description="Daisy SDK control server (FastAPI)")
    parser.add_argument("--host", default=DEFAULT_CONTROL_BIND, help="Bind address for control server (use 0.0.0.0).")
    parser.add_argument("--port", type=int, default=DEFAULT_CONTROL_PORT)
    parser.add_argument("--image", required=True)
    parser.add_argument("--app-port", default=DEFAULT_APP_PORT)
    parser.add_argument("--project-config-path", default=str(DEFAULT_PROJECT_CONFIG_PATH))
    parser.add_argument("--platform", default=None)
    parser.add_argument("--container-name", default=DEFAULT_CONTAINER_NAME)
    parser.add_argument("--detach", action="store_true", help="Run control server in the background.")
    parser.add_argument("--container-logs", action="store_true", help="Stream container logs to stdout.")
    parser.add_argument("--show-docker-command", action="store_true", help="Print docker run command.")
    args = parser.parse_args()

    run_server(
        bind_host=args.host,
        client_host=DEFAULT_CONTROL_HOST,
        port=args.port,
        image=args.image,
        app_port=str(args.app_port),
        project_config_path=Path(args.project_config_path),
        platform=args.platform,
        container_name=args.container_name,
        detach=args.detach,
        log_container=args.container_logs,
        show_docker_command=args.show_docker_command,
    )


if __name__ == "__main__":
    main()
