#!/usr/bin/env python3
"""
Auto-generate markdown docs from docker-compose.yml files.
Usage: python3 scripts/generate_docs.py
"""
import yaml
import os
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).parent.parent
SERVICES_DIR = REPO_ROOT
DOCS_DIR = REPO_ROOT / "docs" / "services"

# Load root .env.example for shared variables
def load_root_env_example():
    env_vars = {}
    env_file = REPO_ROOT / ".env.example"
    if env_file.exists():
        for line in env_file.read_text().splitlines():
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                key, _, val = line.partition("=")
                env_vars[key.strip()] = val.strip()
    return env_vars

# Load service .env.example
def load_service_env_example(service_dir):
    env_vars = {}
    env_file = service_dir / ".env.example"
    if env_file.exists():
        for line in env_file.read_text().splitlines():
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                key, _, val = line.partition("=")
                env_vars[key.strip()] = val.strip()
    return env_vars


def format_volumes(volumes):
    if not volumes:
        return []
    rows = []
    for v in volumes:
        if isinstance(v, str):
            parts = v.split(":")
            if len(parts) >= 2:
                host = parts[0]
                container = parts[1]
                mode = parts[2] if len(parts) > 2 else "rw"
                rows.append((host, container, mode))
            else:
                rows.append((v, "-", "-"))
        elif isinstance(v, dict):
            source = v.get("source", v.get("type", ""))
            target = v.get("target", "")
            mode = v.get("mode", "rw")
            rows.append((source, target, mode))
    return rows


def format_ports(ports):
    if not ports:
        return []
    rows = []
    for p in ports:
        if isinstance(p, str):
            # Handle protocol suffix like "8080:8080/tcp"
            proto = "tcp"
            port_str = p
            if "/" in p:
                port_str, proto = p.rsplit("/", 1)
            # Use rsplit to handle ${VAR:-default}:port correctly
            parts = port_str.rsplit(":", 1)
            if len(parts) == 2:
                rows.append((parts[0], parts[1], proto))
            else:
                rows.append((parts[0], parts[0], proto))
        elif isinstance(p, dict):
            host = str(p.get("published", ""))
            container = str(p.get("target", ""))
            proto = p.get("protocol", "tcp")
            rows.append((host, container, proto))
    return rows


def format_env(env_list):
    if not env_list:
        return []
    rows = []
    for e in env_list:
        if isinstance(e, str) and "=" in e:
            key, _, val = e.partition("=")
            rows.append((key.strip(), val.strip()))
        elif isinstance(e, str):
            rows.append((e.strip(), ""))
    return rows


def format_labels(labels):
    if not labels:
        return []
    rows = []
    if isinstance(labels, list):
        for l in labels:
            if isinstance(l, str) and "=" in l:
                key, _, val = l.partition("=")
                rows.append((key.strip(), val.strip()))
            elif isinstance(l, dict):
                for k, v in l.items():
                    rows.append((k, str(v)))
    elif isinstance(labels, dict):
        for k, v in labels.items():
            rows.append((k, str(v)))
    return rows


def format_command(cmd):
    if not cmd:
        return []
    if isinstance(cmd, str):
        return [cmd]
    if isinstance(cmd, list):
        return [str(c) for c in cmd]
    return []


def format_networks(networks):
    if not networks:
        return []
    rows = []
    if isinstance(networks, dict):
        for name, cfg in networks.items():
            external = False
            if isinstance(cfg, dict):
                external = cfg.get("external", False)
            rows.append((name, "external" if external else "internal"))
    elif isinstance(networks, list):
        for n in networks:
            if isinstance(n, str):
                rows.append((n, "internal"))
    return rows


def format_depends(depends):
    if not depends:
        return []
    if isinstance(depends, list):
        return [str(d) for d in depends]
    if isinstance(depends, dict):
        return list(depends.keys())
    return []


def format_cap_add(caps):
    if not caps:
        return []
    if isinstance(caps, list):
        return [str(c) for c in caps]
    return []


def format_security(sec):
    if not sec:
        return []
    if isinstance(sec, list):
        return [str(s) for s in sec]
    return []


def generate_service_md(service_name, compose_data, root_env, service_env):
    lines = []
    lines.append(f"# {service_name}")
    lines.append("")
    lines.append(f"> Auto-generated from `docker-compose.yml`")
    lines.append("")

    services = compose_data.get("services", {})
    if not services:
        lines.append("No services defined.")
        return "\n".join(lines)

    for svc_name, svc in services.items():
        if len(services) > 1:
            lines.append(f"## `{svc_name}`")
            lines.append("")

        # Image
        image = svc.get("image", "N/A")
        container = svc.get("container_name", svc_name)
        lines.append(f"| Field | Value |")
        lines.append(f"|-------|-------|")
        lines.append(f"| **Image** | `{image}` |")
        lines.append(f"| **Container** | `{container}` |")
        lines.append(f"| **Restart** | `{svc.get('restart', 'no')}` |")
        lines.append("")

        # Command
        cmd = format_command(svc.get("command"))
        if cmd:
            lines.append("### Command")
            lines.append("")
            lines.append("```")
            lines.append(" ".join(cmd))
            lines.append("```")
            lines.append("")

        # Environment
        env = format_env(svc.get("environment"))
        env_file = svc.get("env_file")
        env_from_file = []
        if env_file:
            if isinstance(env_file, str):
                env_file = [env_file]
            env_from_file = env_file

        if env or env_from_file:
            lines.append("### Environment")
            lines.append("")
            if env_from_file:
                lines.append("**Env files:** " + ", ".join(f"`{e}`" for e in env_from_file))
                lines.append("")
            if env:
                lines.append("| Variable | Value |")
                lines.append("|----------|-------|")
                for key, val in env:
                    if "${" in val or val == "":
                        lines.append(f"| `{key}` | `{val}` |")
                    else:
                        lines.append(f"| `{key}` | `{val}` |")
                lines.append("")

        # Service secrets from .env.example
        if service_env:
            lines.append("### Secrets (from `.env.example`)")
            lines.append("")
            lines.append("These variables must be set in the service's `.env` file:")
            lines.append("")
            lines.append("| Variable | Default |")
            lines.append("|----------|---------|")
            for key, val in service_env.items():
                lines.append(f"| `{key}` | `{val}` |")
            lines.append("")

        # Ports
        ports = format_ports(svc.get("ports"))
        if ports:
            lines.append("### Ports")
            lines.append("")
            lines.append("| Host | Container | Protocol |")
            lines.append("|------|-----------|----------|")
            for host, container, proto in ports:
                lines.append(f"| `{host}` | `{container}` | {proto} |")
            lines.append("")

        # Volumes
        volumes = format_volumes(svc.get("volumes"))
        if volumes:
            lines.append("### Volumes")
            lines.append("")
            lines.append("| Host Path | Container Path | Mode |")
            lines.append("|-----------|----------------|------|")
            for host, container, mode in volumes:
                lines.append(f"| `{host}` | `{container}` | {mode} |")
            lines.append("")

        # Networks
        nets = format_networks(svc.get("networks"))
        if nets:
            lines.append("### Networks")
            lines.append("")
            for name, ntype in nets:
                lines.append(f"- `{name}` ({ntype})")
            lines.append("")

        # Depends on
        depends = format_depends(svc.get("depends_on"))
        if depends:
            lines.append("### Dependencies")
            lines.append("")
            for d in depends:
                lines.append(f"- `{d}`")
            lines.append("")

        # Capabilities
        caps = format_cap_add(svc.get("cap_add"))
        if caps:
            lines.append("### Capabilities")
            lines.append("")
            lines.append(", ".join(f"`{c}`" for c in caps))
            lines.append("")

        # Security
        sec = format_security(svc.get("security_opt"))
        if sec:
            lines.append("### Security Options")
            lines.append("")
            for s in sec:
                lines.append(f"- `{s}`")
            lines.append("")

        # Labels
        labels = format_labels(svc.get("labels"))
        if labels:
            lines.append("### Labels")
            lines.append("")
            lines.append("| Key | Value |")
            lines.append("|-----|-------|")
            for key, val in labels:
                lines.append(f"| `{key}` | `{val}` |")
            lines.append("")

        # Devices
        devices = svc.get("devices")
        if devices:
            lines.append("### Devices")
            lines.append("")
            for d in devices:
                if isinstance(d, str):
                    lines.append(f"- `{d}`")
            lines.append("")

        # tmpfs
        tmpfs = svc.get("tmpfs")
        if tmpfs:
            lines.append("### Tmpfs")
            lines.append("")
            if isinstance(tmpfs, list):
                for t in tmpfs:
                    lines.append(f"- `{t}`")
            lines.append("")

        # shm_size
        shm = svc.get("shm_size")
        if shm:
            lines.append(f"### SHM Size: `{shm}`")
            lines.append("")

        # resource limits
        mem = svc.get("mem_limit")
        swap = svc.get("memswap_limit")
        if mem or swap:
            lines.append("### Resources")
            lines.append("")
            if mem:
                lines.append(f"- Memory limit: `{mem}`")
            if swap:
                lines.append(f"- Memory+Swap limit: `{swap}`")
            lines.append("")

        # Compose-level volumes (named)
        comp_volumes = compose_data.get("volumes")
        if comp_volumes and isinstance(comp_volumes, dict):
            lines.append("### Named Volumes (compose-level)")
            lines.append("")
            for vname, vcfg in comp_volumes.items():
                external = False
                if isinstance(vcfg, dict):
                    external = vcfg.get("external", False)
                ntype = "external" if external else "managed by compose"
                lines.append(f"- `{vname}` ({ntype})")
            lines.append("")

        # Compose-level networks
        comp_nets = compose_data.get("networks")
        if comp_nets and isinstance(comp_nets, dict):
            lines.append("### Networks (compose-level)")
            lines.append("")
            for nname, ncfg in comp_nets.items():
                external = False
                if isinstance(ncfg, dict):
                    external = ncfg.get("external", False)
                ntype = "external" if external else "internal"
                lines.append(f"- `{nname}` ({ntype})")
            lines.append("")

        lines.append("---")
        lines.append("")

    return "\n".join(lines)


def generate_index(all_services):
    lines = []
    lines.append("# Homelab Documentation")
    lines.append("")
    lines.append("> Auto-generated from `docker-compose.yml` files")
    lines.append("")

    lines.append("## Data Classification")
    lines.append("")
    lines.append("| Type | Example | Git Repo? | Backup? | Location |")
    lines.append("|------|---------|-----------|---------|----------|")
    lines.append("| **Source Configs** | `docker-compose.yml`, `.env.example`, `Makefile` | Yes | No | `/srv/homelab-infra` |")
    lines.append("| **Service Secrets** | `.env` (DB passwords, API keys) | No | Yes | `/srv/homelab-infra/<service>/.env` |")
    lines.append("| **Runtime Configs** | `acme.json`, service configs | No | Yes | `/mnt/storage/docker-data/<service>` |")
    lines.append("| **Persistent Data** | DB data, uploads, media | No | Yes | `/mnt/storage/docker-data/<service>` |")
    lines.append("")

    lines.append("## Services")
    lines.append("")
    lines.append("| Service | Image | Status |")
    lines.append("|---------|-------|--------|")
    for name, data in sorted(all_services.items()):
        services = data.get("services", {})
        images = []
        for svc in services.values():
            img = svc.get("image", "N/A")
            images.append(f"`{img}`")
        status = "active" if "planned" not in name.lower() else "planned"
        lines.append(f"| [{name}](services/{name}.md) | {', '.join(images)} | {status} |")
    lines.append("")

    lines.append("## Quick Start")
    lines.append("")
    lines.append("```bash")
    lines.append("cp .env.example .env")
    lines.append("for svc in */; do [ -f \"$svc/.env.example\" ] && cp \"$svc/.env.example\" \"$svc/.env\"; done")
    lines.append("docker network create web")
    lines.append("make up")
    lines.append("```")
    lines.append("")

    return "\n".join(lines)


def main():
    root_env = load_root_env_example()
    all_services = {}

    # Find all service directories with docker-compose.yml
    for item in sorted(SERVICES_DIR.iterdir()):
        if not item.is_dir():
            continue
        compose_file = item / "docker-compose.yml"
        if not compose_file.exists():
            continue
        # Skip scripts/docs directories
        if item.name in ("scripts", "docs", "__pycache__"):
            continue

        service_name = item.name
        try:
            with open(compose_file) as f:
                compose_data = yaml.safe_load(f)
        except Exception as e:
            print(f"Error parsing {compose_file}: {e}", file=sys.stderr)
            continue

        if not compose_data:
            continue

        service_env = load_service_env_example(item)
        all_services[service_name] = compose_data

        # Generate service doc
        md = generate_service_md(service_name, compose_data, root_env, service_env)
        out_file = DOCS_DIR / f"{service_name}.md"
        out_file.parent.mkdir(parents=True, exist_ok=True)
        out_file.write_text(md)
        print(f"Generated {out_file.relative_to(REPO_ROOT)}")

    # Generate index
    index_md = generate_index(all_services)
    index_file = REPO_ROOT / "docs" / "index.md"
    index_file.write_text(index_md)
    print(f"Generated {index_file.relative_to(REPO_ROOT)}")

    print(f"\nDone. Generated docs for {len(all_services)} services.")


if __name__ == "__main__":
    main()