"""
aiohttp web server for argonode.

Routes
------
GET  /                  → serve frontend/index.html
GET  /static/{path}     → serve frontend JS/CSS
GET  /nodes             → JSON dict of all registered node definitions
GET  /files             → list files in the current session upload workspace
GET  /folder-files      → list compatible files in a picked folder
GET  /channels          → inspect channels for a picked file
POST /upload            → multipart file upload to the current session workspace
POST /upload-folder     → create a folder in the current session workspace
POST /prompt            → submit a workflow; returns {prompt_id}
GET  /ws                → WebSocket upgrade

WebSocket message types sent to clients
----------------------------------------
{"type": "execution_start",    "data": {"prompt_id": "..."}}
{"type": "executing",          "data": {"node": "...", "prompt_id": "..."}}
{"type": "preview",            "data": {"node_id": "...", "image": "data:..."}}
{"type": "table",              "data": {"node_id": "...", "rows": [...]} }
{"type": "scalar",             "data": {"node_id": "...", "value": 1.23, "unit": "nm"}}
{"type": "node_timing",        "data": {"node_id": "...", "elapsed_ms": 12.34}}
{"type": "execution_error",    "data": {"node_id": "...", "message": "..."}}
{"type": "execution_complete", "data": {"prompt_id": "..."}}
"""

from __future__ import annotations

import asyncio
import json
import logging
import sys
from collections import defaultdict
from copy import deepcopy
from pathlib import Path

from aiohttp import web, WSMsgType

from backend.frontend_build import FrontendBuildError, ensure_frontend_dist_ready
from backend.runtime_paths import ensure_runtime_dirs, frontend_dir, frontend_dist_dir, project_root
from backend.session_runtime import (
    PATH_INPUT_TYPES,
    SESSION_HEADER,
    SESSION_QUERY,
    ensure_session_runtime_dirs,
    normalize_relative_upload_path,
    resolve_client_path,
    server_path_to_client_path,
    session_input_dir,
    session_upload_uri,
    validate_session_id,
)

log = logging.getLogger(__name__)

FRONTEND_DIR = frontend_dir()
DIST_DIR = frontend_dist_dir()
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"


class _SafeEncoder(json.JSONEncoder):
    def default(self, obj):
        import numpy as np

        if isinstance(obj, (np.integer,)):
            return int(obj)
        if isinstance(obj, (np.floating,)):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return super().default(obj)


def _dumps(obj) -> str:
    return json.dumps(obj, cls=_SafeEncoder)


def save_png_bytes(target_path: str, payload: bytes) -> Path:
    path = Path(target_path).expanduser()
    if not target_path.strip():
        raise ValueError("Missing save path")
    if path.suffix.lower() != ".png":
        path = path.with_suffix(".png")
    if not payload.startswith(PNG_SIGNATURE):
        raise ValueError("Payload is not a valid PNG")
    path.write_bytes(payload)
    return path


def create_app(
    loop: asyncio.AbstractEventLoop,
    *,
    allow_local_filesystem: bool = False,
) -> web.Application:
    import backend.nodes  # noqa: F401
    from backend.execution import ExecutionEngine, new_prompt_id
    from backend.node_registry import NODE_CLASS_MAPPINGS, get_all_node_info

    ensure_runtime_dirs()

    session_engines: dict[str, ExecutionEngine] = {}
    session_websockets: dict[str, set[web.WebSocketResponse]] = defaultdict(set)

    def _is_link(value) -> bool:
        return (
            isinstance(value, (list, tuple))
            and len(value) == 2
            and isinstance(value[0], str)
            and isinstance(value[1], int)
        )

    def require_session_id(request: web.Request) -> str:
        raw_session = request.headers.get(SESSION_HEADER) or request.query.get(SESSION_QUERY)
        if not raw_session:
            if allow_local_filesystem:
                raw_session = "desktop-local-session"
            else:
                raise web.HTTPBadRequest(reason="Missing session id")

        try:
            session_id = validate_session_id(raw_session)
        except ValueError as exc:
            raise web.HTTPBadRequest(reason=str(exc)) from exc

        ensure_session_runtime_dirs(session_id)
        return session_id

    def get_session_engine(session_id: str) -> ExecutionEngine:
        engine = session_engines.get(session_id)
        if engine is None:
            engine = ExecutionEngine()
            session_engines[session_id] = engine
        return engine

    def resolve_request_path(session_id: str, raw_value: str) -> Path:
        try:
            return resolve_client_path(
                raw_value,
                session_id=session_id,
                allow_local_filesystem=allow_local_filesystem,
            )
        except PermissionError as exc:
            raise web.HTTPForbidden(reason=str(exc)) from exc
        except ValueError as exc:
            raise web.HTTPBadRequest(reason=str(exc)) from exc

    def rewrite_prompt_paths(prompt: dict, session_id: str) -> dict:
        normalized = deepcopy(prompt)
        for node_def in normalized.values():
            class_name = node_def.get("class_type")
            cls = NODE_CLASS_MAPPINGS.get(class_name)
            if cls is None:
                continue

            input_types = cls.INPUT_TYPES()
            specs = {}
            specs.update(input_types.get("required", {}))
            specs.update(input_types.get("optional", {}))

            inputs = node_def.get("inputs", {})
            if not isinstance(inputs, dict):
                continue

            for input_name, raw_value in list(inputs.items()):
                if _is_link(raw_value) or not isinstance(raw_value, str):
                    continue
                if not raw_value.strip():
                    continue

                spec = specs.get(input_name)
                input_type = spec[0] if isinstance(spec, (list, tuple)) and spec else spec
                if not isinstance(input_type, str):
                    continue
                if input_type not in PATH_INPUT_TYPES:
                    continue

                inputs[input_name] = str(resolve_request_path(session_id, raw_value))
        return normalized

    def broadcast(session_id: str, msg: dict) -> None:
        payload = _dumps(msg)
        for ws in list(session_websockets.get(session_id, ())):
            if not ws.closed:
                asyncio.run_coroutine_threadsafe(ws.send_str(payload), loop)

    def on_preview(session_id: str, node_id: str, data_uri: str) -> None:
        broadcast(session_id, {"type": "preview", "data": {"node_id": node_id, "image": data_uri}})

    def on_table(session_id: str, node_id: str, rows: list) -> None:
        broadcast(session_id, {"type": "table", "data": {"node_id": node_id, "rows": rows}})

    def on_mesh(session_id: str, node_id: str, mesh_data: dict) -> None:
        broadcast(session_id, {"type": "mesh3d", "data": {"node_id": node_id, "mesh": mesh_data}})

    def on_overlay(session_id: str, node_id: str, overlay_data) -> None:
        broadcast(session_id, {"type": "overlay", "data": {"node_id": node_id, "overlay": overlay_data}})

    def on_value(session_id: str, node_id: str, payload) -> None:
        if isinstance(payload, dict):
            value = payload.get("value")
            unit = payload.get("unit", "")
        else:
            value = payload
            unit = ""

        data = {"node_id": node_id, "value": value}
        if isinstance(unit, str) and unit.strip():
            data["unit"] = unit.strip()
        broadcast(session_id, {"type": "scalar", "data": data})

    def on_warning(session_id: str, node_id: str, message: str) -> None:
        broadcast(session_id, {"type": "node_warning", "data": {"node_id": node_id, "message": message}})

    async def index(request: web.Request) -> web.Response:
        if not getattr(sys, "frozen", False):
            try:
                await loop.run_in_executor(
                    None,
                    lambda: ensure_frontend_dist_ready(
                        project_root(),
                        FRONTEND_DIR,
                        DIST_DIR,
                        logger=log,
                    ),
                )
            except FrontendBuildError as exc:
                log.error("Unable to refresh frontend build: %s", exc)
                return web.Response(status=500, text=str(exc), content_type="text/plain")

        if (DIST_DIR / "index.html").exists():
            return web.FileResponse(DIST_DIR / "index.html")
        return web.Response(
            status=500,
            text=(
                "Frontend build not found. Run `npm run build` from the repo root, "
                "or use `npm run dev` for the Vite development server."
            ),
            content_type="text/plain",
        )

    async def get_nodes(request: web.Request) -> web.Response:
        return web.Response(
            text=_dumps(get_all_node_info()),
            content_type="application/json",
        )

    async def list_files(request: web.Request) -> web.Response:
        session_id = require_session_id(request)
        input_path = session_input_dir(session_id)
        files = sorted(
            server_path_to_client_path(entry, session_id)
            for entry in input_path.iterdir()
            if entry.is_file() and not entry.name.startswith(".")
        ) if input_path.exists() else []
        return web.Response(text=_dumps(files), content_type="application/json")

    async def create_upload_folder(request: web.Request) -> web.Response:
        session_id = require_session_id(request)
        body = await request.json()
        relative_path = normalize_relative_upload_path(body.get("path", ""))
        target = session_input_dir(session_id) / Path(relative_path.as_posix())
        target.mkdir(parents=True, exist_ok=True)
        return web.Response(
            text=_dumps({"path": session_upload_uri(relative_path)}),
            content_type="application/json",
        )

    async def get_folder_files(request: web.Request) -> web.Response:
        from backend.nodes.helpers import list_folder_paths

        session_id = require_session_id(request)
        folder_path = request.query.get("folder", "")
        if not folder_path:
            return web.Response(text=_dumps([]), content_type="application/json")

        resolved_path = resolve_request_path(session_id, folder_path)
        running_loop = asyncio.get_running_loop()
        entries = await running_loop.run_in_executor(None, list_folder_paths, str(resolved_path))

        payload = []
        for entry in entries:
            mapped = dict(entry)
            if "path" in mapped:
                mapped["path"] = server_path_to_client_path(mapped["path"], session_id)
            payload.append(mapped)
        return web.Response(text=_dumps(payload), content_type="application/json")

    async def upload_file(request: web.Request) -> web.Response:
        session_id = require_session_id(request)
        reader = await request.multipart()
        relative_path = None
        filename = ""
        file_bytes = None

        while True:
            field = await reader.next()
            if field is None:
                break
            if field.name == "relative_path":
                relative_path = await field.text()
                continue
            if field.name == "file":
                filename = Path(field.filename or "upload.bin").name
                chunks = []
                while True:
                    chunk = await field.read_chunk(65536)
                    if not chunk:
                        break
                    chunks.append(chunk)
                file_bytes = b"".join(chunks)

        if file_bytes is None:
            raise web.HTTPBadRequest(reason="Expected a 'file' field in multipart body")

        relative = normalize_relative_upload_path(relative_path or filename)
        dest = session_input_dir(session_id) / Path(relative.as_posix())
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(file_bytes)

        return web.Response(
            text=_dumps({"filename": filename, "path": session_upload_uri(relative)}),
            content_type="application/json",
        )

    async def download_file(request: web.Request) -> web.Response:
        body = await request.read()
        filename = request.query.get("filename", "workflow.png")
        return web.Response(
            body=body,
            content_type="application/octet-stream",
            headers={"Content-Disposition": f'attachment; filename="{filename}"'},
        )

    async def save_workflow_png(request: web.Request) -> web.Response:
        body = await request.read()
        target_path = request.query.get("path", "")
        if not target_path:
            raise web.HTTPBadRequest(reason="Missing path")
        try:
            saved_path = save_png_bytes(target_path, body)
        except ValueError as exc:
            raise web.HTTPBadRequest(reason=str(exc)) from exc
        return web.Response(
            text=_dumps({"path": str(saved_path)}),
            content_type="application/json",
        )

    async def get_channels(request: web.Request) -> web.Response:
        from backend.nodes.helpers import list_channels

        session_id = require_session_id(request)
        filepath = request.query.get("file", "")
        if not filepath:
            return web.Response(
                text=_dumps([{"name": "field", "type": "DATA_FIELD"}]),
                content_type="application/json",
            )

        resolved_path = resolve_request_path(session_id, filepath)
        channels = await loop.run_in_executor(None, list_channels, str(resolved_path))
        return web.Response(text=_dumps(channels), content_type="application/json")

    async def submit_prompt(request: web.Request) -> web.Response:
        session_id = require_session_id(request)
        body = await request.json()
        prompt = body.get("prompt")
        if not isinstance(prompt, dict) or not prompt:
            raise web.HTTPBadRequest(reason="'prompt' must be a non-empty dict")

        normalized_prompt = rewrite_prompt_paths(prompt, session_id)
        prompt_id = new_prompt_id()
        engine = get_session_engine(session_id)

        async def run():
            broadcast(session_id, {"type": "execution_start", "data": {"prompt_id": prompt_id}})

            def on_start(node_id: str) -> None:
                broadcast(session_id, {"type": "executing", "data": {"node": node_id, "prompt_id": prompt_id}})

            def on_done(node_id: str, elapsed_ms: float) -> None:
                broadcast(session_id, {
                    "type": "node_timing",
                    "data": {"node_id": node_id, "elapsed_ms": elapsed_ms},
                })

            try:
                await loop.run_in_executor(
                    None,
                    lambda: engine.execute(
                        normalized_prompt,
                        on_node_start=on_start,
                        on_node_done=on_done,
                        on_preview=lambda node_id, payload: on_preview(session_id, node_id, payload),
                        on_table=lambda node_id, rows: on_table(session_id, node_id, rows),
                        on_mesh=lambda node_id, mesh_data: on_mesh(session_id, node_id, mesh_data),
                        on_overlay=lambda node_id, overlay_data: on_overlay(session_id, node_id, overlay_data),
                        on_value=lambda node_id, payload: on_value(session_id, node_id, payload),
                        on_warning=lambda node_id, message: on_warning(session_id, node_id, message),
                    ),
                )
                broadcast(session_id, {"type": "execution_complete", "data": {"prompt_id": prompt_id}})
            except Exception as exc:
                log.exception("Execution error")
                broadcast(session_id, {
                    "type": "execution_error",
                    "data": {"node_id": "", "message": str(exc)},
                })

        asyncio.ensure_future(run())
        return web.Response(
            text=_dumps({"prompt_id": prompt_id}),
            content_type="application/json",
        )

    async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
        session_id = require_session_id(request)
        ws = web.WebSocketResponse()
        await ws.prepare(request)
        session_websockets[session_id].add(ws)
        log.info(
            "WebSocket client connected for session %s (%d total in session)",
            session_id,
            len(session_websockets[session_id]),
        )
        try:
            async for msg in ws:
                if msg.type == WSMsgType.TEXT:
                    pass
                elif msg.type in (WSMsgType.ERROR, WSMsgType.CLOSE):
                    break
        finally:
            session_websockets[session_id].discard(ws)
            if not session_websockets[session_id]:
                session_websockets.pop(session_id, None)
            log.info(
                "WebSocket client disconnected for session %s (%d remaining in session)",
                session_id,
                len(session_websockets.get(session_id, ())),
            )
        return ws

    app = web.Application()
    app["allow_local_filesystem"] = allow_local_filesystem

    app.router.add_get("/", index)
    app.router.add_get("/nodes", get_nodes)
    app.router.add_get("/files", list_files)
    app.router.add_get("/folder-files", get_folder_files)
    app.router.add_post("/upload-folder", create_upload_folder)
    app.router.add_post("/upload", upload_file)
    app.router.add_post("/download", download_file)
    app.router.add_post("/save-workflow-png", save_workflow_png)
    app.router.add_get("/channels", get_channels)
    app.router.add_post("/prompt", submit_prompt)
    app.router.add_get("/ws", websocket_handler)

    if (DIST_DIR / "assets").exists():
        app.router.add_static("/assets", DIST_DIR / "assets")
    if FRONTEND_DIR.exists():
        app.router.add_static("/static", FRONTEND_DIR)

    async def _cors_middleware(app_, handler):
        async def middleware(request):
            if request.method == "OPTIONS":
                return web.Response(headers={
                    "Access-Control-Allow-Origin": "*",
                    "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
                    "Access-Control-Allow-Headers": f"Content-Type, {SESSION_HEADER}",
                })
            response = await handler(request)
            response.headers["Access-Control-Allow-Origin"] = "*"
            return response

        return middleware

    app.middlewares.append(_cors_middleware)
    return app