#!/usr/bin/env python3

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from tempfile import NamedTemporaryFile
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError

from codex_ws_auths import enable_codex_auth_websockets
from xai_auth_proxy import enforce_config_proxy_url, parse_proxy_urls, stable_proxy_for_auth


def env_required(name: str) -> str:
    value = os.environ.get(name, "").strip()
    if not value:
        raise SystemExit(f"missing required env: {name}")
    return value


MC_BIN = os.environ.get("MC_BIN", "mc")
ALIAS = os.environ.get("OBJECTSTORE_ALIAS", "daili-objectstore").strip() or "daili-objectstore"
ENDPOINT = env_required("OBJECTSTORE_ENDPOINT")
ACCESS_KEY = env_required("OBJECTSTORE_ACCESS_KEY")
SECRET_KEY = env_required("OBJECTSTORE_SECRET_KEY")
BUCKET = env_required("OBJECTSTORE_BUCKET")
ROOT = Path(env_required("OBJECTSTORE_ROOT")).resolve()
CONFIG_FALLBACK = os.environ.get("OBJECTSTORE_CONFIG_FALLBACK", "").strip()
REMOTE_CONFIG_KEY = "config/config.yaml"
REMOTE_AUTHS_PREFIX = "auths"
STATE_PATH = Path(os.environ.get("OBJECTSTORE_SYNC_STATE_PATH", "") or (ROOT / ".objectstore-sync-state.json"))
XAI_PROXY_URLS = parse_proxy_urls(
    os.environ.get("XAI_PROXY_URLS", "") or "socks5h://127.0.0.1:9"
)


def mc_env() -> dict:
    env = os.environ.copy()
    for key in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"):
        env.pop(key, None)
    return env


MC_TIMEOUT = int(os.environ.get("MC_TIMEOUT", "30"))


def run_mc(args: list[str], check: bool = True) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        [MC_BIN, *args],
        check=check,
        env=mc_env(),
        capture_output=True,
        text=True,
        timeout=MC_TIMEOUT,
    )


def ensure_alias() -> None:
    run_mc(["alias", "set", ALIAS, ENDPOINT, ACCESS_KEY, SECRET_KEY, "--api", "S3v4", "--path", "on"])


def remote_path(key: str) -> str:
    return f"{ALIAS}/{BUCKET}/{key}"


def parse_ts(value: str) -> float:
    return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp()


def local_path(rel: str) -> Path:
    return ROOT / rel


def _is_auth_rel(rel: str) -> bool:
    return rel.startswith(f"{REMOTE_AUTHS_PREFIX}/") and len(rel) > len(REMOTE_AUTHS_PREFIX) + 1


def load_sync_state() -> tuple[set[str], dict[str, dict[str, str]]]:
    try:
        data = json.loads(STATE_PATH.read_text())
    except (OSError, ValueError):
        return set(), {}
    if not isinstance(data, dict):
        return set(), {}
    raw_rels = data.get("known_rels")
    known_rels = {str(rel) for rel in raw_rels if isinstance(rel, str)} if isinstance(raw_rels, list) else set()
    derived: dict[str, dict[str, str]] = {}
    raw_derived = data.get("derived_from_remote")
    if isinstance(raw_derived, dict):
        for rel, value in raw_derived.items():
            if not isinstance(rel, str) or not isinstance(value, dict):
                continue
            remote_etag = value.get("remote_etag")
            local_md5 = value.get("local_md5")
            if isinstance(remote_etag, str) and isinstance(local_md5, str):
                derived[rel] = {
                    "remote_etag": remote_etag,
                    "local_md5": local_md5,
                }
    return known_rels, derived


def load_known_rels() -> set[str]:
    return load_sync_state()[0]


def save_known_rels(
    inventory: dict[str, dict],
    derived_from_remote: dict[str, dict[str, str]] | None = None,
) -> None:
    derived = derived_from_remote or {}
    payload = {
        "version": 2,
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "known_rels": sorted(inventory),
        "derived_from_remote": {
            rel: derived[rel]
            for rel in sorted(derived)
            if rel in inventory
        },
    }
    STATE_PATH.parent.mkdir(parents=True, exist_ok=True)
    with NamedTemporaryFile(
        "w",
        encoding="utf-8",
        delete=False,
        dir=str(STATE_PATH.parent),
    ) as tmp:
        json.dump(payload, tmp, ensure_ascii=False, sort_keys=True)
        tmp.write("\n")
        tmp_path = Path(tmp.name)
    tmp_path.replace(STATE_PATH)


def touch_local(path: Path, ts: float) -> None:
    os.utime(path, (ts, ts))


def prune_empty_dirs(root: Path) -> None:
    if not root.exists():
        return
    for path in sorted((p for p in root.rglob("*") if p.is_dir()), reverse=True):
        try:
            path.rmdir()
        except OSError:
            pass


def file_md5(path: Path) -> str:
    digest = hashlib.md5()
    with path.open("rb") as handle:
        while True:
            chunk = handle.read(1024 * 1024)
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()


def remote_inventory() -> dict[str, dict]:
    inventory: dict[str, dict] = {}

    config_proc = run_mc(["stat", "--json", remote_path(REMOTE_CONFIG_KEY)], check=False)
    if config_proc.returncode == 0 and config_proc.stdout.strip():
        raw = json.loads(config_proc.stdout)
        inventory[REMOTE_CONFIG_KEY] = {
            "etag": raw["etag"],
            "last_modified": parse_ts(raw["lastModified"]),
            "remote": remote_path(REMOTE_CONFIG_KEY),
        }

    auths_proc = run_mc(["ls", "--json", "--recursive", remote_path(REMOTE_AUTHS_PREFIX)], check=False)
    if auths_proc.returncode == 0:
        for line in auths_proc.stdout.splitlines():
            if not line.strip():
                continue
            raw = json.loads(line)
            if raw.get("status") != "success" or raw.get("type") != "file":
                continue
            rel = f"{REMOTE_AUTHS_PREFIX}/{raw['key']}"
            inventory[rel] = {
                "etag": raw["etag"],
                "last_modified": parse_ts(raw["lastModified"]),
                "remote": remote_path(rel),
            }

    return inventory


def local_inventory() -> dict[str, dict]:
    inventory: dict[str, dict] = {}

    config_path = local_path(REMOTE_CONFIG_KEY)
    if config_path.is_file():
        inventory[REMOTE_CONFIG_KEY] = {
            "md5": file_md5(config_path),
            "mtime": config_path.stat().st_mtime,
        }

    auth_root = local_path(REMOTE_AUTHS_PREFIX)
    if auth_root.is_dir():
        for path in sorted(auth_root.rglob("*")):
            if not path.is_file():
                continue
            rel = path.relative_to(ROOT).as_posix()
            inventory[rel] = {
                "md5": file_md5(path),
                "mtime": path.stat().st_mtime,
            }

    return inventory


def _enforce_downloaded_xai_proxy(path: Path, auth_name: str) -> bool:
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError):
        return False
    if not isinstance(payload, dict) or str(payload.get("type") or "").strip().lower() != "xai":
        return False

    desired_proxy = stable_proxy_for_auth(auth_name, XAI_PROXY_URLS)
    current_proxy = payload.get("proxy_url")
    current_proxy = current_proxy.strip() if isinstance(current_proxy, str) else ""
    changed = current_proxy != desired_proxy
    if changed:
        payload["proxy_url"] = desired_proxy
        path.write_text(
            json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
    path.chmod(0o600)
    return changed


def download_file(rel: str, meta: dict) -> bool:
    dest = local_path(rel)
    dest.parent.mkdir(parents=True, exist_ok=True)
    with NamedTemporaryFile(delete=False, dir=str(dest.parent)) as temporary:
        temporary_path = Path(temporary.name)
    temporary_path.unlink()
    replaced = False
    try:
        run_mc(["cp", meta["remote"], str(temporary_path)])
        changed = False
        if rel == REMOTE_CONFIG_KEY:
            changed = enforce_config_proxy_url(temporary_path)
        elif _is_auth_rel(rel):
            changed = _enforce_downloaded_xai_proxy(temporary_path, dest.name)
        touch_local(temporary_path, meta["last_modified"])
        temporary_path.replace(dest)
        replaced = True
    finally:
        if not replaced and temporary_path.exists():
            temporary_path.unlink()
    return changed


def remember_derived_file(
    derived_from_remote: dict[str, dict[str, str]],
    rel: str,
    remote_etag: str,
) -> None:
    path = local_path(rel)
    derived_from_remote[rel] = {
        "remote_etag": remote_etag,
        "local_md5": file_md5(path),
    }


def enforce_config_without_mtime_change(path: Path) -> bool:
    original_mtime = path.stat().st_mtime
    changed = enforce_config_proxy_url(path)
    if changed:
        touch_local(path, original_mtime)
    return changed


def upload_file(rel: str) -> None:
    src = local_path(rel)
    if not src.is_file():
        return
    run_mc(["cp", str(src), remote_path(rel)])


def delete_remote(rel: str) -> None:
    run_mc(["rm", "--force", remote_path(rel)], check=False)


INVALID_STATUS_KEYWORDS = [
    "token_invalidated",
    "token_revoked",
    "invalidated oauth token",
    "authentication token has been invalidated",
    "account has been deactivated",
    "no_organization",
    "unauthorized",
    "401",
]


def _is_free_auth_file_name(file_name: str) -> bool:
    return file_name.strip().endswith("-free.json")


def _free_auth_cleanup_enabled() -> bool:
    return os.environ.get("CLEANUP_ENABLED", "0").strip() == "1"


def _fetch_invalid_auth_names() -> set[str]:
    """Query local CLIProxyAPI management API for invalid free auth names.

    Team auth lifecycle is managed externally; runtime sync must not evict team
    auths for transient states such as cooldown or temporary unavailability.
    """
    if not _free_auth_cleanup_enabled():
        return set()

    mgmt_key = os.environ.get("MANAGEMENT_PASSWORD") or os.environ.get("API_KEY") or ""
    port = os.environ.get("PORT", "8317")
    if not mgmt_key:
        return set()
    url = f"http://127.0.0.1:{port}/v0/management/auth-files"
    req = Request(url)
    req.add_header("Authorization", f"Bearer {mgmt_key}")
    try:
        with urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read())
    except (HTTPError, URLError, OSError, ValueError):
        return set()
    if not isinstance(data, dict):
        return set()
    files = data.get("files", [])
    if not isinstance(files, list):
        return set()

    invalid_names: set[str] = set()
    now = datetime.now(timezone.utc).isoformat()
    for entry in files:
        if not isinstance(entry, dict):
            continue
        name = str(entry.get("name") or "").strip()
        if not name:
            continue
        if not _is_free_auth_file_name(name):
            continue
        status = str(entry.get("status") or "").strip().lower()
        status_message = str(entry.get("status_message") or "").strip().lower()
        unavailable = entry.get("unavailable", False)
        # Check next_retry_after — if still in cooldown, treat as invalid
        next_retry = str(entry.get("next_retry_after") or "").strip()
        in_cooldown = bool(next_retry and next_retry > now)
        if status == "error" or unavailable or in_cooldown:
            invalid_names.add(name)
            continue
        if status_message:
            for kw in INVALID_STATUS_KEYWORDS:
                if kw in status_message:
                    invalid_names.add(name)
                    break
    return invalid_names


def restore() -> None:
    ensure_alias()
    ROOT.mkdir(parents=True, exist_ok=True)
    local_path("config").mkdir(parents=True, exist_ok=True)
    shutil.rmtree(local_path(REMOTE_AUTHS_PREFIX), ignore_errors=True)
    local_path(REMOTE_AUTHS_PREFIX).mkdir(parents=True, exist_ok=True)

    remote = remote_inventory()
    derived_from_remote: dict[str, dict[str, str]] = {}
    if REMOTE_CONFIG_KEY in remote:
        if download_file(REMOTE_CONFIG_KEY, remote[REMOTE_CONFIG_KEY]):
            remember_derived_file(
                derived_from_remote,
                REMOTE_CONFIG_KEY,
                remote[REMOTE_CONFIG_KEY]["etag"],
            )
    elif CONFIG_FALLBACK:
        fallback = Path(CONFIG_FALLBACK)
        if fallback.is_file():
            target = local_path(REMOTE_CONFIG_KEY)
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(fallback, target)

    run_mc(["cp", "--recursive", remote_path(f"{REMOTE_AUTHS_PREFIX}/"), str(local_path(REMOTE_AUTHS_PREFIX))], check=False)
    for rel, meta in remote.items():
        if not _is_auth_rel(rel):
            continue
        path = local_path(rel)
        if not path.is_file():
            continue
        changed = _enforce_downloaded_xai_proxy(path, path.name)
        touch_local(path, meta["last_modified"])
        if changed:
            remember_derived_file(derived_from_remote, rel, meta["etag"])
    enable_codex_auth_websockets(local_path(REMOTE_AUTHS_PREFIX))
    config_path = local_path(REMOTE_CONFIG_KEY)
    if config_path.is_file():
        enforce_config_without_mtime_change(config_path)
    save_known_rels(local_inventory(), derived_from_remote)


def sync() -> None:
    ensure_alias()
    ROOT.mkdir(parents=True, exist_ok=True)
    enable_codex_auth_websockets(local_path(REMOTE_AUTHS_PREFIX))
    config_path = local_path(REMOTE_CONFIG_KEY)
    if config_path.is_file():
        enforce_config_without_mtime_change(config_path)

    invalid_names = _fetch_invalid_auth_names()
    known_rels, derived_from_remote = load_sync_state()

    remote = remote_inventory()
    local = local_inventory()

    # Remote has, local doesn't:
    #   - if this auth file was seen in a previous sync, local deletion wins
    #     and we delete the remote mirror too. This makes management-UI DELETE
    #     durable instead of letting the next sync resurrect the file.
    #   - if it is genuinely new to this instance, download it.
    # Both have, md5 differs → compare mtime, newer wins (unless invalid)
    for rel, meta in remote.items():
        # Extract auth file name from rel path (e.g. "auths/codex-xxx-free.json" → "codex-xxx-free.json")
        file_name = rel.rsplit("/", 1)[-1] if "/" in rel else rel
        if file_name in invalid_names:
            derived_from_remote.pop(rel, None)
            delete_remote(rel)
            dest = local_path(rel)
            if dest.is_file():
                dest.unlink()
            continue

        local_meta = local.get(rel)
        if local_meta is None:
            if _is_auth_rel(rel) and rel in known_rels:
                derived_from_remote.pop(rel, None)
                delete_remote(rel)
                continue
            if download_file(rel, meta):
                remember_derived_file(derived_from_remote, rel, meta["etag"])
            else:
                derived_from_remote.pop(rel, None)
            continue
        if local_meta["md5"] == meta["etag"]:
            derived_from_remote.pop(rel, None)
            continue
        derived = derived_from_remote.get(rel)
        if (
            derived is not None
            and derived.get("remote_etag") == meta["etag"]
            and derived.get("local_md5") == local_meta["md5"]
        ):
            continue
        if local_meta["mtime"] > meta["last_modified"]:
            derived_from_remote.pop(rel, None)
            upload_file(rel)
        else:
            if download_file(rel, meta):
                remember_derived_file(derived_from_remote, rel, meta["etag"])
            else:
                derived_from_remote.pop(rel, None)

    # Local has, remote doesn't → upload to remote (new file from management UI)
    for rel in sorted(set(local) - set(remote)):
        file_name = rel.rsplit("/", 1)[-1] if "/" in rel else rel
        if file_name in invalid_names:
            derived_from_remote.pop(rel, None)
            dest = local_path(rel)
            if dest.is_file():
                dest.unlink()
            continue
        derived_from_remote.pop(rel, None)
        upload_file(rel)

    if config_path.is_file():
        enforce_config_without_mtime_change(config_path)
    prune_empty_dirs(local_path(REMOTE_AUTHS_PREFIX))
    save_known_rels(local_inventory(), derived_from_remote)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", choices=("restore", "sync"))
    args = parser.parse_args()

    if args.mode == "restore":
        restore()
    else:
        sync()
    return 0


if __name__ == "__main__":
    sys.exit(main())