#!/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

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"))

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]] = {}
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 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) -> bytes:
    headers = [
        "HTTP/1.1 302 Found",
        f"Location: {location}",
        "Content-Length: 0",
        "Connection: close",
        "",
        "",
    ]
    return "\r\n".join(headers).encode("utf-8")


def build_user_host(host: str, user_id: str) -> Optional[str]:
    if not host:
        return None

    hostname, separator, port = host.partition(":")
    hostname = hostname.strip().lower()
    if not hostname:
        return None

    if hostname == "localhost":
        target_host = f"{user_id}.localhost"
    else:
        labels = [label for label in hostname.split(".") if label]
        if not labels:
            return None
        if labels[0] == user_id:
            target_host = hostname
        else:
            target_host = ".".join([user_id, *labels])

    return f"{target_host}:{port}" if separator else target_host


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

    header_end = data.find(b"\r\n\r\n")
    if header_end == -1:
        return None, data, 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, data, 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()

    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

    path, separator, query = target.partition("?")
    segments = [segment for segment in path.split("/") if segment]
    if segments:
        candidate = segments[0]
        if candidate in known_users:
            user_host = build_user_host(headers.get("host", ""), candidate)
            if user_host:
                redirect_path = "/" + "/".join(segments[1:]) if len(segments) > 1 else "/"
                if separator:
                    redirect_path = f"{redirect_path}?{query}"
                redirect_target = f"http://{user_host}{redirect_path}"
                return None, data, redirect_target

    return None, data, None


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")
    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

    user, rewritten_data, redirect_target = extract_user_from_http_request(initial_data, users_by_id)
    if redirect_target:
        writer.write(build_http_redirect(redirect_target))
        await writer.drain()
        writer.close()
        await writer.wait_closed()
        return

    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 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(f"Shared gateway ready on port {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"])
        log(f"Ready on port {user['port']} for user {user['id']} ({user['container_name']})")
        servers.append(server)

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


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