#!/usr/bin/env python3
"""
Container manager for Fiscality sessions.

The manager exposes one TCP listener per configured user. On first incoming
connection it creates or starts that user's dedicated Fiscality container and
proxies traffic to the internal application port.
"""
import asyncio
import json
import os
import socket
import threading
import time
from typing import Any, Dict, Iterable, Optional
from urllib.parse import urlsplit

import docker

CONFIG_PATH = os.getenv("CONFIG_PATH", "config.json")
IMAGE_NAME = os.getenv("IMAGE_NAME", "fis")
APP_PORT = int(os.getenv("APP_PORT", "3000"))
GATEWAY_PORT = int(os.getenv("GATEWAY_PORT", "0"))
INACTIVITY_SECONDS = int(os.getenv("INACTIVITY_SECONDS", "10"))
IDLE_CHECK_INTERVAL = int(os.getenv("IDLE_CHECK_INTERVAL", "10"))
STARTUP_TIMEOUT_SECONDS = int(os.getenv("STARTUP_TIMEOUT_SECONDS", "60"))
STARTUP_STABILIZATION_SECONDS = float(os.getenv("STARTUP_STABILIZATION_SECONDS", "5"))
TCP_CONNECT_TIMEOUT_SECONDS = float(os.getenv("TCP_CONNECT_TIMEOUT_SECONDS", "1.5"))
INITIAL_REQUEST_TIMEOUT_SECONDS = float(os.getenv("INITIAL_REQUEST_TIMEOUT_SECONDS", "3"))
MAX_REQUEST_HEADER_BYTES = int(os.getenv("MAX_REQUEST_HEADER_BYTES", "65536"))
GATEWAY_ROUTE_TTL_SECONDS = int(os.getenv("GATEWAY_ROUTE_TTL_SECONDS", "3600"))
GATEWAY_ROUTE_COOKIE_NAME = "fis_user_route"

DEFAULTS: Dict[str, Any] = {
    "container_prefix": "fis-",
    "environment": {},
    "volumes": {},
    "mem_limit": None,
    "memswap_limit": None,
    "privileged": False,
    "tmpfs": {},
    "dns": [],
    "extra_hosts": {},
}

_docker_client = None
user_activity: Dict[str, float] = {}
users_by_id: Dict[str, Dict[str, Any]] = {}
gateway_client_routes: Dict[str, Dict[str, Any]] = {}
user_container_locks: Dict[str, threading.Lock] = {}
user_container_locks_guard = threading.Lock()


def get_docker_client():
    global _docker_client
    if _docker_client is None:
        _docker_client = docker.from_env()
    return _docker_client


def log(message: str) -> None:
    print(message, flush=True)


def format_access_url(host: str, port: int) -> str:
    return f"http://{host}:{port}"


def get_user_container_lock(user_id: str) -> threading.Lock:
    with user_container_locks_guard:
        if user_id not in user_container_locks:
            user_container_locks[user_id] = threading.Lock()
        return user_container_locks[user_id]


def merge_dict(base: Dict[str, Any], override: Optional[Dict[str, Any]]) -> Dict[str, Any]:
    result = dict(base or {})
    result.update(override or {})
    return result


def resolve_placeholders(value: Any, variables: Dict[str, Any]) -> Any:
    if isinstance(value, str):
        return value.format(**variables)
    if isinstance(value, dict):
        return {
            resolve_placeholders(key, variables): resolve_placeholders(item, variables)
            for key, item in value.items()
        }
    if isinstance(value, list):
        return [resolve_placeholders(item, variables) for item in value]
    return value


def merge_defaults(config: Dict[str, Any], user: Dict[str, Any]) -> Dict[str, Any]:
    defaults = merge_dict(DEFAULTS, config.get("defaults"))
    merged = dict(user)

    merged["id"] = str(user["id"])
    merged["port"] = int(user["port"])
    variables = {"id": merged["id"], "port": merged["port"]}

    merged["environment"] = merge_dict(
        resolve_placeholders(defaults.get("environment", {}), variables),
        resolve_placeholders(user.get("environment"), variables),
    )
    merged["volumes"] = merge_dict(
        resolve_placeholders(defaults.get("volumes", {}), variables),
        resolve_placeholders(user.get("volumes"), variables),
    )
    merged["extra_hosts"] = merge_dict(
        resolve_placeholders(defaults.get("extra_hosts", {}), variables),
        resolve_placeholders(user.get("extra_hosts"), variables),
    )
    merged["tmpfs"] = merge_dict(
        resolve_placeholders(defaults.get("tmpfs", {}), variables),
        resolve_placeholders(user.get("tmpfs"), variables),
    )
    merged["dns"] = resolve_placeholders(user.get("dns", defaults.get("dns", [])), variables)
    merged["mem_limit"] = user.get("mem_limit", defaults.get("mem_limit"))
    merged["memswap_limit"] = user.get("memswap_limit", defaults.get("memswap_limit"))
    merged["privileged"] = user.get("privileged", defaults.get("privileged", False))

    if not merged.get("container_name"):
        prefix = defaults.get("container_prefix", "fis-")
        merged["container_name"] = f"{prefix}{merged['id']}"

    if not merged["volumes"]:
        merged["volumes"] = {f"v-conf-{merged['id']}": "/config"}

    return merged


def validate_users(users: Iterable[Dict[str, Any]]) -> None:
    seen_ids = set()
    seen_ports = set()
    seen_names = set()

    for user in users:
        uid = user["id"]
        port = user["port"]
        name = user["container_name"]

        if uid in seen_ids:
            raise ValueError(f"Duplicate user id: {uid}")
        if port in seen_ports:
            raise ValueError(f"Duplicate port: {port}")
        if name in seen_names:
            raise ValueError(f"Duplicate container_name: {name}")
        if not user["volumes"]:
            raise ValueError(f"User {uid} must define at least one volume")
        if "PASSWORD" not in user["environment"] or not user["environment"]["PASSWORD"]:
            raise ValueError(f"User {uid} must define environment.PASSWORD")

        seen_ids.add(uid)
        seen_ports.add(port)
        seen_names.add(name)


def load_config(path: str = CONFIG_PATH) -> Dict[str, Any]:
    with open(path, encoding="utf-8") as handle:
        raw = json.load(handle)

    raw_users = raw.get("users", [])
    if not raw_users:
        raise ValueError("Configuration must contain at least one user")

    users = [merge_defaults(raw, user) for user in raw_users]
    validate_users(users)

    return {"defaults": merge_dict(DEFAULTS, raw.get("defaults")), "users": users}


def get_container_ip(container_name: str) -> Optional[str]:
    """Get IP address of a running container."""
    try:
        container = get_docker_client().containers.get(container_name)
        container.reload()
        if container.status != "running":
            return None
        networks = container.attrs["NetworkSettings"]["Networks"]
        for network in networks.values():
            if network.get("IPAddress"):
                return network["IPAddress"]
    except docker.errors.NotFound:
        return None
    except Exception as exc:
        log(f"[{container_name}] Failed to read container IP: {exc}")
    return None


def is_port_open(host: str, port: int, timeout_seconds: float = TCP_CONNECT_TIMEOUT_SECONDS) -> bool:
    try:
        with socket.create_connection((host, port), timeout=timeout_seconds):
            return True
    except OSError:
        return False


async def read_initial_http_request(reader) -> bytes:
    buffer = bytearray()
    while len(buffer) < MAX_REQUEST_HEADER_BYTES:
        chunk = await asyncio.wait_for(reader.read(4096), timeout=INITIAL_REQUEST_TIMEOUT_SECONDS)
        if not chunk:
            break
        buffer.extend(chunk)
        if b"\r\n\r\n" in buffer:
            break
    return bytes(buffer)


def build_http_error(status_code: int, reason: str, message: str) -> bytes:
    body = f"{message}\n".encode("utf-8")
    headers = [
        f"HTTP/1.1 {status_code} {reason}",
        "Content-Type: text/plain; charset=utf-8",
        f"Content-Length: {len(body)}",
        "Connection: close",
        "",
        "",
    ]
    return "\r\n".join(headers).encode("utf-8") + body


def build_http_redirect(location: str, extra_headers: Optional[Iterable[str]] = None) -> bytes:
    headers = [
        "HTTP/1.1 302 Found",
        f"Location: {location}",
        "Content-Length: 0",
        "Connection: close",
    ]
    headers.extend(extra_headers or [])
    headers.extend(["", ""])
    return "\r\n".join(headers).encode("utf-8")


def rewrite_http_request_target(data: bytes, target: str) -> bytes:
    header_end = data.find(b"\r\n")
    if header_end == -1:
        return data

    request_line = data[:header_end].decode("iso-8859-1", errors="replace")
    parts = request_line.split(" ", 2)
    if len(parts) != 3:
        return data

    method, _, version = parts
    rewritten_line = f"{method} {target} {version}".encode("iso-8859-1")
    return rewritten_line + data[header_end:]


def parse_http_request(data: bytes):
    if not data:
        return None

    header_end = data.find(b"\r\n\r\n")
    if header_end == -1:
        return None

    header_block = data[:header_end].decode("iso-8859-1", errors="replace")
    lines = header_block.split("\r\n")
    if not lines or len(lines[0].split(" ")) < 3:
        return None

    method, target, version = lines[0].split(" ", 2)
    headers = {}
    for line in lines[1:]:
        if ":" not in line:
            continue
        key, value = line.split(":", 1)
        headers[key.strip().lower()] = value.strip()

    return {
        "method": method,
        "target": target,
        "version": version,
        "headers": headers,
    }


def extract_cookie_user(headers: Dict[str, str], known_users: Dict[str, Dict[str, Any]]) -> Optional[Dict[str, Any]]:
    cookie_header = headers.get("cookie", "")
    if not cookie_header:
        return None

    for part in cookie_header.split(";"):
        key, separator, value = part.strip().partition("=")
        if separator and key == GATEWAY_ROUTE_COOKIE_NAME:
            user_id = value.strip()
            return known_users.get(user_id)
    return None


def is_websocket_request(headers: Dict[str, str]) -> bool:
    return headers.get("upgrade", "").lower() == "websocket"


def extract_request_target(data: bytes) -> Optional[str]:
    parsed = parse_http_request(data)
    if not parsed:
        return None
    return parsed["target"]


def remember_gateway_route(client_ip: str, user_id: str) -> None:
    if not client_ip:
        return
    gateway_client_routes[client_ip] = {"user_id": user_id, "updated_at": time.time()}


def get_remembered_gateway_user(client_ip: str) -> Optional[Dict[str, Any]]:
    if not client_ip:
        return None

    route = gateway_client_routes.get(client_ip)
    if not route:
        return None

    if time.time() - route["updated_at"] > GATEWAY_ROUTE_TTL_SECONDS:
        gateway_client_routes.pop(client_ip, None)
        return None

    user_id = route["user_id"]
    return users_by_id.get(user_id)


def extract_user_from_http_request(data: bytes, known_users: Dict[str, Dict[str, Any]]):
    parsed = parse_http_request(data)
    if not parsed:
        return None, data, None, False

    target = parsed["target"]
    headers = parsed["headers"]

    host = headers.get("host", "").split(":", 1)[0].strip().lower()
    if host:
        labels = [label for label in host.split(".") if label]
        if labels:
            candidate = labels[0]
            if candidate in known_users and candidate not in {"localhost", "127", "0"}:
                return known_users[candidate], data, None, False

    parsed_target = urlsplit(target)
    if parsed_target.scheme and parsed_target.netloc:
        path = parsed_target.path or "/"
        query = parsed_target.query
        separator = "?" if query else ""
    else:
        path, separator, query = target.partition("?")

    segments = [segment for segment in path.split("/") if segment]
    if segments:
        candidate = segments[0]
        if candidate in known_users:
            rewritten_path = "/" + "/".join(segments[1:]) if len(segments) > 1 else "/"
            rewritten_target = f"{rewritten_path}?{query}" if separator else rewritten_path
            rewritten_data = rewrite_http_request_target(data, rewritten_target)
            return known_users[candidate], rewritten_data, None, True

    cookie_user = extract_cookie_user(headers, known_users)
    if cookie_user:
        return cookie_user, data, None, False

    return None, data, None, False


def wait_until_ready(container_name: str) -> bool:
    stabilized = False
    deadline = time.time() + STARTUP_TIMEOUT_SECONDS
    while time.time() < deadline:
        ip = get_container_ip(container_name)
        if ip and is_port_open(ip, APP_PORT):
            if STARTUP_STABILIZATION_SECONDS > 0 and not stabilized:
                log(
                    f"[{container_name}] Port {APP_PORT} is open, waiting "
                    f"{STARTUP_STABILIZATION_SECONDS}s for backend stabilization"
                )
                time.sleep(STARTUP_STABILIZATION_SECONDS)
                stabilized = True
                continue
            return True
        time.sleep(0.5)
    return False


def ensure_docker_volumes(volume_bindings: Dict[str, str]) -> Dict[str, Dict[str, str]]:
    client = get_docker_client()
    volumes = {}
    for volume_name, mount_path in volume_bindings.items():
        try:
            client.volumes.get(volume_name)
        except docker.errors.NotFound:
            client.volumes.create(name=volume_name)
        volumes[volume_name] = {"bind": mount_path, "mode": "rw"}
    return volumes


def start_user_container(user: Dict[str, Any]):
    """Start or create the requested user's container."""
    name = user["container_name"]
    client = get_docker_client()
    lock = get_user_container_lock(user["id"])

    with lock:
        try:
            container = client.containers.get(name)
            container.reload()

            if container.status == "running":
                log(f"[{user['id']}] Container already running")
                return container

            log(f"[{user['id']}] Starting existing container {name}")
            container.start()
        except docker.errors.NotFound:
            log(f"[{user['id']}] Creating container {name}")
            try:
                container = client.containers.run(
                    IMAGE_NAME,
                    name=name,
                    detach=True,
                    volumes=ensure_docker_volumes(user["volumes"]),
                    environment=user.get("environment", {}),
                    mem_limit=user.get("mem_limit"),
                    memswap_limit=user.get("memswap_limit"),
                    privileged=user.get("privileged", False),
                    tmpfs=user.get("tmpfs") or None,
                    dns=user.get("dns") or None,
                    extra_hosts=user.get("extra_hosts") or None,
                    tty=True,
                    stdin_open=True,
                    restart_policy={"Name": "no"},
                )
            except docker.errors.APIError as exc:
                if "already in use" not in str(exc).lower():
                    raise
                log(f"[{user['id']}] Container {name} was created concurrently, reusing it")
                container = client.containers.get(name)
                container.reload()
                if container.status != "running":
                    container.start()

    if wait_until_ready(name):
        log(f"[{user['id']}] Ready")
    else:
        log(f"[{user['id']}] Startup timeout after {STARTUP_TIMEOUT_SECONDS}s")

    return container


def stop_user_container(user: Dict[str, Any]) -> None:
    name = user["container_name"]
    try:
        container = get_docker_client().containers.get(name)
    except docker.errors.NotFound:
        return

    try:
        container.stop(timeout=10)
        log(f"[{user['id']}] Stopped container {name}")
    except Exception as exc:
        log(f"[{user['id']}] Failed to stop container {name}: {exc}")


async def proxy_data(reader, writer, backend_ip: str, backend_port: int, user_id: str, initial_data: bytes = b""):
    """Simple bidirectional proxy."""
    try:
        backend_reader, backend_writer = await asyncio.open_connection(backend_ip, backend_port)
    except Exception as exc:
        log(f"[{user_id}] Can't connect to backend {backend_ip}:{backend_port}: {exc}")
        writer.close()
        await writer.wait_closed()
        return

    if initial_data:
        backend_writer.write(initial_data)
        await backend_writer.drain()
        user_activity[user_id] = time.time()

    async def pipe(src, dst):
        try:
            while True:
                data = await src.read(8192)
                if not data:
                    break
                dst.write(data)
                await dst.drain()
                user_activity[user_id] = time.time()
        except Exception:
            pass

    tasks = [
        asyncio.create_task(pipe(reader, backend_writer)),
        asyncio.create_task(pipe(backend_reader, writer)),
    ]

    await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)

    for task in tasks:
        task.cancel()

    try:
        backend_writer.close()
        writer.close()
        await backend_writer.wait_closed()
        await writer.wait_closed()
    except Exception:
        pass


async def handle_connection(reader, writer, user: Dict[str, Any], initial_data: bytes = b""):
    """Handle incoming connection for a configured user."""
    user_id = user["id"]
    log(f"[{user_id}] New connection")

    loop = asyncio.get_running_loop()
    try:
        await loop.run_in_executor(None, start_user_container, user)
    except Exception as exc:
        log(f"[{user_id}] Failed to prepare container: {exc}")
        writer.close()
        await writer.wait_closed()
        return

    ip = get_container_ip(user["container_name"])
    if not ip:
        log(f"[{user_id}] No container IP available")
        writer.close()
        await writer.wait_closed()
        return

    user_activity[user_id] = time.time()
    await proxy_data(reader, writer, ip, APP_PORT, user_id, initial_data=initial_data)


async def handle_gateway_connection(reader, writer):
    peer = writer.get_extra_info("peername")
    client_ip = peer[0] if isinstance(peer, tuple) and peer else None
    try:
        initial_data = await read_initial_http_request(reader)
    except asyncio.TimeoutError:
        writer.write(build_http_error(408, "Request Timeout", "Gateway request timeout"))
        await writer.drain()
        writer.close()
        await writer.wait_closed()
        return

    parsed_request = parse_http_request(initial_data)
    user, rewritten_data, redirect_target, selected_via_path = extract_user_from_http_request(initial_data, users_by_id)
    if user:
        remember_gateway_route(client_ip, user["id"])

    if redirect_target:
        writer.write(build_http_redirect(redirect_target))
        await writer.drain()
        writer.close()
        await writer.wait_closed()
        return

    if user and selected_via_path and parsed_request and not is_websocket_request(parsed_request["headers"]):
        normalized_target = extract_request_target(rewritten_data) or "/"
        writer.write(
            build_http_redirect(
                normalized_target,
                extra_headers=[
                    f"Set-Cookie: {GATEWAY_ROUTE_COOKIE_NAME}={user['id']}; Path=/; HttpOnly; SameSite=Lax"
                ],
            )
        )
        await writer.drain()
        writer.close()
        await writer.wait_closed()
        log(f"[gateway] Switched {client_ip} to user {user['id']} via path route")
        return

    if not user:
        remembered_user = get_remembered_gateway_user(client_ip)
        if remembered_user:
            user = remembered_user
            rewritten_data = initial_data
            log(f"[gateway] Reused remembered route for {client_ip} -> user {user['id']}")

    if not user:
        writer.write(
            build_http_error(
                404,
                "Not Found",
                "Unknown Fiscality user route. Use /<user-id> or a matching Host header.",
            )
        )
        await writer.drain()
        writer.close()
        await writer.wait_closed()
        return

    log(f"[gateway] Routed {peer} to user {user['id']}")
    await handle_connection(reader, writer, user, initial_data=rewritten_data)


async def idle_checker():
    """Stop idle containers using the configured container_name."""
    while True:
        await asyncio.sleep(IDLE_CHECK_INTERVAL)
        now = time.time()

        for client_ip, route in list(gateway_client_routes.items()):
            if now - route["updated_at"] > GATEWAY_ROUTE_TTL_SECONDS:
                gateway_client_routes.pop(client_ip, None)

        for user_id, last_active in list(user_activity.items()):
            if now - last_active <= INACTIVITY_SECONDS:
                continue

            user = users_by_id.get(user_id)
            if not user:
                user_activity.pop(user_id, None)
                continue

            log(f"[{user_id}] Idle timeout reached, stopping container")
            stop_user_container(user)
            user_activity.pop(user_id, None)


async def main():
    config = load_config(CONFIG_PATH)
    global users_by_id
    users_by_id = {user["id"]: user for user in config["users"]}

    log("Starting FIS manager")
    log(f"Config path: {CONFIG_PATH}")
    log(f"Image name: {IMAGE_NAME}")
    log(f"Application port: {APP_PORT}")
    log(f"Gateway port: {GATEWAY_PORT if GATEWAY_PORT else 'disabled'}")
    log(f"Inactivity timeout: {INACTIVITY_SECONDS}s")

    servers = []
    if GATEWAY_PORT > 0:
        gateway_server = await asyncio.start_server(handle_gateway_connection, "0.0.0.0", GATEWAY_PORT)
        log(
            "Shared gateway ready: "
            f"{format_access_url('localhost', GATEWAY_PORT)} "
            f"(host-based routing: http://<user>.localhost:{GATEWAY_PORT})"
        )
        servers.append(gateway_server)

    for user in config["users"]:
        async def handler(reader, writer, configured_user=user):
            await handle_connection(reader, writer, configured_user)

        server = await asyncio.start_server(handler, "0.0.0.0", user["port"])
        direct_url = format_access_url("localhost", user["port"])
        if GATEWAY_PORT > 0:
            gateway_url = format_access_url(f"{user['id']}.localhost", GATEWAY_PORT)
            log(
                f"Ready for user {user['id']} ({user['container_name']}): "
                f"{direct_url} or {gateway_url}"
            )
        else:
            log(f"Ready for user {user['id']} ({user['container_name']}): {direct_url}")
        servers.append(server)

    asyncio.create_task(idle_checker())
    await asyncio.gather(*(server.serve_forever() for server in servers))


if __name__ == "__main__":
    asyncio.run(main())