"""Same-origin HTTP transport for localhost and Hugging Face Docker Spaces."""

import argparse
import hashlib
import json
import logging
import os
import re
import threading
from http.cookies import SimpleCookie
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import unquote, urlparse

from game.logging_config import configure_logging
from game.media import _is_displayable_image
from server.sessions import Conflict, Sessions
from server.views import catalog

ROOT = Path(__file__).resolve().parents[1]
MAX_BODY = 5 * 1024 * 1024


class Assets:
    def __init__(self):
        self.paths = {}
        self.lock = threading.Lock()
        self.roots = [
            ROOT / "data" / name
            for name in ("static_location_images", "static_npc_images", "ui")
        ]
        self.roots += [ROOT / "images_cache"]
        # Discover the actual portrait folder without exposing the whole data tree.
        from game.vlm_engine import STATIC_NPC_IMAGES_DIR

        self.roots.append(Path(STATIC_NPC_IMAGES_DIR).resolve())

    def url(self, value):
        if not isinstance(value, str):
            return None
        path = Path(value)
        path = (ROOT / path).resolve() if not path.is_absolute() else path.resolve()
        if not any(
            path.is_relative_to(root.resolve()) for root in self.roots
        ) or not _is_displayable_image(path):
            return None
        key = hashlib.sha256(str(path).encode()).hexdigest()
        with self.lock:
            self.paths[key] = path
        return "/media/" + key


def make_handler(sessions, assets, space_host=None):
    # Trust only deployment configuration, never forwarded headers from requests.
    if space_host and not re.fullmatch(
        r"[a-z0-9]+(?:-[a-z0-9]+)*\.hf\.space", space_host
    ):
        raise ValueError("SPACE_HOST must be a Hugging Face Space hostname.")
    public_origin = "https://" + space_host if space_host else None
    ancestors = "'self' https://huggingface.co" if space_host else "'none'"
    cookie_name = "__Host-ww_session" if space_host else "ww_session"
    cookie_policy = (
        "SameSite=None; Secure; Partitioned" if space_host else "SameSite=Strict"
    )

    class Handler(BaseHTTPRequestHandler):
        server_version = "WandsWhispers/1"

        def log_message(self, fmt, *args):
            logging.getLogger("http").debug(fmt, *args)

        def headers_for(self, status, content_type, length):
            self.send_response(status)
            self.send_header("Content-Type", content_type)
            self.send_header("Content-Length", str(length))
            self.send_header("Cache-Control", "no-store")
            self.send_header("X-Content-Type-Options", "nosniff")
            self.send_header("Referrer-Policy", "same-origin")
            self.send_header(
                "Content-Security-Policy",
                f"default-src 'self'; script-src 'self'; style-src 'self' https://fonts.googleapis.com; font-src 'self' https://fonts.gstatic.com; img-src 'self' blob: data:; connect-src 'self'; frame-ancestors {ancestors}; base-uri 'self'; form-action 'self'",
            )
            if getattr(self, "new_cookie", None):
                self.send_header(
                    "Set-Cookie",
                    f"{cookie_name}={self.new_cookie}; HttpOnly; {cookie_policy}; Path=/; Max-Age=86400",
                )
            self.end_headers()

        def respond(self, status, value):
            body = json.dumps(value, ensure_ascii=False, allow_nan=False).encode()
            self.headers_for(status, "application/json; charset=utf-8", len(body))
            self.wfile.write(body)

        def session(self):
            cookie = SimpleCookie()
            try:
                cookie.load(self.headers.get("Cookie", ""))
            except Exception:
                pass
            supplied = cookie.get(cookie_name)
            old = supplied.value if supplied else None
            token, session = sessions.get(old)
            if token != old:
                self.new_cookie = token
            return session

        def valid_origin(self):
            host = self.headers.get("Host", "")
            local_hosts = {
                f"127.0.0.1:{self.server.server_port}",
                f"localhost:{self.server.server_port}",
            }
            if host not in local_hosts and host != space_host:
                return False
            expected = public_origin if host == space_host else "http://" + host
            origin = self.headers.get("Origin")
            # The Hub embeds the initial HTML cross-site; its API calls still
            # originate from the exact hf.space origin. Never exempt API routes.
            embedded_document = (
                bool(space_host)
                and self.command == "GET"
                and urlparse(self.path).path
                in {"/", "/web", "/web/", "/web/index.html"}
                and self.headers.get("Sec-Fetch-Dest")
                in {"iframe", "frame", "document"}
                and origin in {None, "https://huggingface.co", public_origin}
            )
            if embedded_document:
                return True
            return (not origin or origin == expected) and self.headers.get(
                "Sec-Fetch-Site"
            ) != "cross-site"

        def do_GET(self):
            try:
                if not self.valid_origin():
                    return self.respond(403, {"error": "Same-origin access required."})
                path = unquote(urlparse(self.path).path)
                if path == "/api/catalog":
                    return self.respond(200, catalog())
                if path == "/api/session":
                    return self.respond(
                        200, sessions.snapshot(self.session(), assets.url)
                    )
                if path == "/api/save":
                    return self.respond(200, sessions.save(self.session()))
                if path.startswith("/media/"):
                    with assets.lock:
                        file = assets.paths.get(path.removeprefix("/media/"))
                    if file is None or assets.url(str(file)) is None:
                        return self.respond(404, {"error": "Image not found."})
                elif path.startswith("/data/static_location_images/"):
                    # Main menu uses the existing castle when its LFS image exists.
                    file = (ROOT / path.lstrip("/")).resolve()
                    if assets.url(str(file)) is None:
                        return self.respond(404, {"error": "Image not available."})
                else:
                    relative = path.removeprefix("/web/").lstrip("/")
                    if path in ("/", "/web", "/web/"):
                        relative = "index.html"
                    file = (ROOT / "web" / relative).resolve()
                    if not file.is_relative_to(ROOT / "web") or file.suffix not in (
                        ".html",
                        ".css",
                        ".js",
                        ".svg",
                        ".jpg",
                        ".png",
                    ):
                        return self.respond(404, {"error": "Not found."})
                if not file.is_file():
                    return self.respond(404, {"error": "Not found."})
                import mimetypes

                body = file.read_bytes()
                self.headers_for(
                    200,
                    mimetypes.guess_type(file)[0] or "application/octet-stream",
                    len(body),
                )
                self.wfile.write(body)
            except Conflict as exc:
                self.respond(409, {"error": str(exc)})
            except ValueError as exc:
                self.respond(400, {"error": str(exc)})
            except (BrokenPipeError, ConnectionResetError):
                pass
            except Exception:
                logging.getLogger("http").exception("GET failed")
                self.respond(
                    500, {"error": "The server could not complete the request."}
                )

        def do_POST(self):
            try:
                if not self.valid_origin() or self.headers.get("X-WW-Client") != "web":
                    return self.respond(403, {"error": "Same-origin access required."})
                if (
                    self.headers.get("Content-Type", "").split(";")[0]
                    != "application/json"
                ):
                    return self.respond(415, {"error": "Send JSON."})
                length = int(self.headers.get("Content-Length", "0"))
                if not 0 < length <= MAX_BODY:
                    return self.respond(
                        413, {"error": "Request must be between 1 byte and 5 MB."}
                    )
                data = json.loads(
                    self.rfile.read(length),
                    parse_constant=lambda _: (_ for _ in ()).throw(
                        ValueError("Invalid number.")
                    ),
                )
                if not isinstance(data, dict):
                    raise ValueError("Expected a JSON object.")
                session = self.session()
                path = urlparse(self.path).path
                if path == "/api/new":
                    sessions.replace(session, data)
                elif path == "/api/load":
                    sessions.replace(session, data, load=True)
                elif path == "/api/settings":
                    sessions.configure(session, data)
                elif path == "/api/stats":
                    sessions.allocate(session, data)
                elif path == "/api/turn":
                    sessions.start_turn(session, data)
                else:
                    return self.respond(404, {"error": "Unknown API endpoint."})
                return self.respond(200, sessions.snapshot(session, assets.url))
            except Conflict as exc:
                self.respond(409, {"error": str(exc)})
            except (ValueError, TypeError, KeyError) as exc:
                self.respond(400, {"error": str(exc)})
            except (BrokenPipeError, ConnectionResetError):
                pass
            except Exception:
                logging.getLogger("http").exception("POST failed")
                self.respond(
                    500, {"error": "The server could not complete the request."}
                )

    return Handler


def main():
    parser = argparse.ArgumentParser(
        description="Wands & Whispers custom UI and game API"
    )
    parser.add_argument("--port", type=int, default=8081)
    parser.add_argument("--host", default="127.0.0.1")
    args = parser.parse_args()
    configure_logging()
    text = image = None
    token = os.environ.get("HF_TOKEN")
    if token:
        from huggingface_hub import InferenceClient
        from openai import OpenAI

        text = OpenAI(
            base_url="https://router.huggingface.co/v1",
            api_key=token,
            timeout=45,
            max_retries=1,
        )
        image = InferenceClient(provider="fal-ai", api_key=token, timeout=45)
    prompt = (ROOT / "prompt/default_system.txt").read_text()
    sessions = Sessions(text, image, prompt)
    server = ThreadingHTTPServer(
        (args.host, args.port),
        make_handler(sessions, Assets(), space_host=os.environ.get("SPACE_HOST")),
    )
    print(
        f"Wands & Whispers: http://127.0.0.1:{args.port}/web/ — AI {'ready' if token else 'requires HF_TOKEN'}",
        flush=True,
    )
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()
        sessions.pool.shutdown(wait=False, cancel_futures=True)


if __name__ == "__main__":
    main()