from __future__ import annotations

import json
import os
import re
import shutil
import tempfile
from dataclasses import dataclass
from datetime import datetime, timezone
from getpass import getpass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import urlparse

import tomlkit


AGENT_CHOICES = ("claude", "codex", "all")
CLAUDE_PROFILE_BASENAME = "settings"
CODEX_PROFILE_BASENAME = "auth"
DEFAULT_IMPORT_SEPARATOR = " - "


class AgswitchError(RuntimeError):
    pass


@dataclass
class ClaudeProfile:
    auth_token: str
    base_url: str
    model: Optional[str] = None
    disable_nonessential_traffic: Optional[str] = "1"

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ClaudeProfile":
        return cls(
            auth_token=require_text(data, "auth_token"),
            base_url=require_text(data, "base_url"),
            model=optional_text(data, "model"),
            disable_nonessential_traffic=optional_text(
                data, "disable_nonessential_traffic"
            ),
        )

    def to_dict(self) -> Dict[str, Any]:
        data: Dict[str, Any] = {
            "auth_token": self.auth_token,
            "base_url": self.base_url,
        }
        if self.model:
            data["model"] = self.model
        if self.disable_nonessential_traffic is not None:
            data["disable_nonessential_traffic"] = self.disable_nonessential_traffic
        return data


@dataclass
class CodexProfile:
    api_key: str
    base_url: str
    provider_name: str
    wire_api: str = "responses"
    requires_openai_auth: bool = True
    auth_mode: str = "apikey"

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "CodexProfile":
        return cls(
            api_key=require_text(data, "api_key"),
            base_url=require_text(data, "base_url"),
            provider_name=require_text(data, "provider_name"),
            wire_api=optional_text(data, "wire_api") or "responses",
            requires_openai_auth=bool(data.get("requires_openai_auth", True)),
            auth_mode=optional_text(data, "auth_mode") or "apikey",
        )

    def to_dict(self) -> Dict[str, Any]:
        return {
            "api_key": self.api_key,
            "base_url": self.base_url,
            "provider_name": self.provider_name,
            "wire_api": self.wire_api,
            "requires_openai_auth": self.requires_openai_auth,
            "auth_mode": self.auth_mode,
        }


@dataclass
class Profile:
    name: str
    claude: Optional[ClaudeProfile] = None
    codex: Optional[CodexProfile] = None

    @classmethod
    def from_dict(cls, name: str, data: Dict[str, Any]) -> "Profile":
        return cls(
            name=name,
            claude=ClaudeProfile.from_dict(data["claude"]) if data.get("claude") else None,
            codex=CodexProfile.from_dict(data["codex"]) if data.get("codex") else None,
        )

    def to_dict(self) -> Dict[str, Any]:
        data: Dict[str, Any] = {}
        if self.claude:
            data["claude"] = self.claude.to_dict()
        if self.codex:
            data["codex"] = self.codex.to_dict()
        return data


@dataclass
class AppPaths:
    home: Path
    store_path: Path
    backup_root: Path
    claude_settings: Path
    codex_auth: Path
    codex_config: Path

    @classmethod
    def discover(cls) -> "AppPaths":
        home = resolve_home()
        store_path = resolve_store_path(home)
        backup_root = store_path.parent / "backups"
        return cls(
            home=home,
            store_path=store_path,
            backup_root=backup_root,
            claude_settings=home / ".claude" / "settings.json",
            codex_auth=home / ".codex" / "auth.json",
            codex_config=home / ".codex" / "config.toml",
        )


@dataclass
class LiveStatus:
    agent: str
    profile_name: Optional[str]
    details: Dict[str, Any]
    paths: List[Path]


def resolve_home() -> Path:
    override = os.environ.get("AGSWITCH_HOME")
    if override:
        return Path(override).expanduser().resolve()
    return Path.home().resolve()


def resolve_store_path(home: Path) -> Path:
    override = os.environ.get("AGSWITCH_STORE")
    if override:
        return Path(override).expanduser().resolve()
    xdg_config = os.environ.get("XDG_CONFIG_HOME")
    if xdg_config:
        return Path(xdg_config).expanduser().resolve() / "agentkeyswitch" / "profiles.json"
    return home / ".config" / "agentkeyswitch" / "profiles.json"


def require_text(data: Dict[str, Any], key: str) -> str:
    value = optional_text(data, key)
    if not value:
        raise AgswitchError(f"Missing required field: {key}")
    return value


def optional_text(data: Dict[str, Any], key: str) -> Optional[str]:
    value = data.get(key)
    if value is None:
        return None
    text = str(value).strip()
    return text or None


def mask_secret(value: Optional[str]) -> str:
    if not value:
        return "-"
    if len(value) <= 8:
        return "*" * len(value)
    return f"{value[:4]}...{value[-4:]}"


def slug_from_url(url: str, fallback: str = "provider") -> str:
    parsed = urlparse(url)
    host = parsed.netloc.lower()
    host = re.sub(r"^www\.", "", host)
    if host:
        candidate = host.split(".")[0]
    else:
        candidate = fallback
    candidate = re.sub(r"[^a-z0-9_]+", "_", candidate)
    candidate = candidate.strip("_") or fallback
    return candidate


def ensure_parent(path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)


def read_json_file(path: Path, missing_ok: bool = False) -> Dict[str, Any]:
    if not path.exists():
        if missing_ok:
            return {}
        raise AgswitchError(f"File not found: {path}")
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        raise AgswitchError(f"Invalid JSON in {path}: {exc}") from exc


def atomic_write_text(path: Path, text: str) -> None:
    ensure_parent(path)
    mode = path.stat().st_mode & 0o777 if path.exists() else 0o600
    fd, temp_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    temp_path = Path(temp_name)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            handle.write(text)
            handle.flush()
            os.fsync(handle.fileno())
        os.chmod(temp_path, mode)
        os.replace(temp_path, path)
    finally:
        if temp_path.exists():
            temp_path.unlink()


def atomic_write_json(path: Path, data: Dict[str, Any]) -> None:
    atomic_write_text(path, json.dumps(data, indent=2, ensure_ascii=True) + "\n")


def backup_file(path: Path, backup_root: Path) -> Optional[Path]:
    if not path.exists():
        return None
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    backup_dir = backup_root / timestamp
    backup_dir.mkdir(parents=True, exist_ok=True)
    target = backup_dir / path.name
    suffix = 1
    while target.exists():
        target = backup_dir / f"{path.name}.{suffix}"
        suffix += 1
    shutil.copy2(path, target)
    return target


def load_profiles(app_paths: AppPaths) -> Dict[str, Profile]:
    if not app_paths.store_path.exists():
        return {}
    raw = read_json_file(app_paths.store_path)
    version = raw.get("version", 1)
    if version != 1:
        raise AgswitchError(
            f"Unsupported profile store version {version} in {app_paths.store_path}"
        )
    profiles = raw.get("profiles", {})
    if not isinstance(profiles, dict):
        raise AgswitchError(f"Invalid profile store format in {app_paths.store_path}")
    return {name: Profile.from_dict(name, data) for name, data in profiles.items()}


def save_profiles(app_paths: AppPaths, profiles: Dict[str, Profile]) -> None:
    ensure_parent(app_paths.store_path)
    try:
        app_paths.store_path.parent.chmod(0o700)
    except FileNotFoundError:
        pass
    data = {
        "version": 1,
        "profiles": {name: profiles[name].to_dict() for name in sorted(profiles)},
    }
    atomic_write_json(app_paths.store_path, data)
    os.chmod(app_paths.store_path, 0o600)


def list_profiles(app_paths: AppPaths) -> List[Profile]:
    return [profile for _, profile in sorted(load_profiles(app_paths).items())]


def build_profile_from_inputs(
    name: str,
    agent: str,
    *,
    claude_token: Optional[str] = None,
    claude_base_url: Optional[str] = None,
    claude_model: Optional[str] = None,
    codex_key: Optional[str] = None,
    codex_base_url: Optional[str] = None,
    codex_provider: Optional[str] = None,
    codex_wire_api: str = "responses",
) -> Profile:
    if agent not in AGENT_CHOICES:
        raise AgswitchError(f"Unsupported agent: {agent}")

    profile = Profile(name=name)
    if agent in ("claude", "all"):
        token = claude_token or prompt_secret("Claude auth token")
        base_url = claude_base_url or prompt_text("Claude base URL")
        model = claude_model if claude_model is not None else prompt_text(
            "Claude model", default="opus"
        )
        profile.claude = ClaudeProfile(
            auth_token=token,
            base_url=base_url,
            model=model,
        )
    if agent in ("codex", "all"):
        key = codex_key or prompt_secret("Codex API key")
        base_url = codex_base_url or prompt_text("Codex base URL")
        provider = codex_provider or prompt_text(
            "Codex provider name", default=slug_from_url(base_url)
        )
        wire_api = codex_wire_api or "responses"
        profile.codex = CodexProfile(
            api_key=key,
            base_url=base_url,
            provider_name=provider,
            wire_api=wire_api,
        )
    return profile


def prompt_text(label: str, default: Optional[str] = None) -> str:
    suffix = f" [{default}]" if default else ""
    value = input(f"{label}{suffix}: ").strip()
    if value:
        return value
    if default is not None:
        return default
    raise AgswitchError(f"{label} is required")


def prompt_secret(label: str) -> str:
    value = getpass(f"{label}: ").strip()
    if not value:
        raise AgswitchError(f"{label} is required")
    return value


def merge_profile(
    existing: Optional[Profile],
    incoming: Profile,
    *,
    overwrite: bool,
) -> Profile:
    result = Profile(name=incoming.name)
    if existing:
        result.claude = existing.claude
        result.codex = existing.codex
    if incoming.claude:
        if result.claude and not overwrite:
            raise AgswitchError(
                f"Profile '{incoming.name}' already has a claude section. Use --overwrite."
            )
        result.claude = incoming.claude
    if incoming.codex:
        if result.codex and not overwrite:
            raise AgswitchError(
                f"Profile '{incoming.name}' already has a codex section. Use --overwrite."
            )
        result.codex = incoming.codex
    return result


def add_profile(
    app_paths: AppPaths,
    incoming: Profile,
    *,
    overwrite: bool,
) -> Profile:
    profiles = load_profiles(app_paths)
    merged = merge_profile(profiles.get(incoming.name), incoming, overwrite=overwrite)
    profiles[incoming.name] = merged
    save_profiles(app_paths, profiles)
    return merged


def remove_profile(app_paths: AppPaths, name: str) -> None:
    profiles = load_profiles(app_paths)
    if name not in profiles:
        raise AgswitchError(f"Unknown profile: {name}")
    del profiles[name]
    save_profiles(app_paths, profiles)


def parse_claude_settings(path: Path) -> ClaudeProfile:
    data = read_json_file(path)
    env = data.get("env")
    if not isinstance(env, dict):
        raise AgswitchError(f"{path} does not contain an 'env' object")
    token = optional_text(env, "ANTHROPIC_AUTH_TOKEN")
    base_url = optional_text(env, "ANTHROPIC_BASE_URL")
    if not token or not base_url:
        raise AgswitchError(
            f"{path} is missing ANTHROPIC_AUTH_TOKEN or ANTHROPIC_BASE_URL"
        )
    return ClaudeProfile(
        auth_token=token,
        base_url=base_url,
        model=optional_text(data, "model"),
        disable_nonessential_traffic=optional_text(
            env, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC"
        ),
    )


def parse_codex_files(
    auth_path: Path,
    config_path: Path,
    *,
    default_provider_name: Optional[str] = None,
) -> CodexProfile:
    auth_data = read_json_file(auth_path)
    api_key = optional_text(auth_data, "OPENAI_API_KEY")
    auth_mode = optional_text(auth_data, "auth_mode") or "apikey"
    auth_base_url = optional_text(auth_data, "base_url")
    if not api_key:
        raise AgswitchError(f"{auth_path} is missing OPENAI_API_KEY")

    provider_name: Optional[str] = None
    base_url = auth_base_url
    wire_api = "responses"
    requires_openai_auth = True
    config_data: Any = {}

    if config_path.exists():
        try:
            config_data = tomlkit.parse(config_path.read_text(encoding="utf-8"))
        except Exception as exc:
            raise AgswitchError(f"Invalid TOML in {config_path}: {exc}") from exc
        current_provider = optional_text(config_data, "model_provider")
        providers = config_data.get("model_providers") or {}
        if current_provider and current_provider in providers:
            provider = providers[current_provider]
            provider_base_url = optional_text(provider, "base_url")
            if not base_url or provider_base_url == base_url:
                provider_name = current_provider
                base_url = base_url or provider_base_url
                wire_api = optional_text(provider, "wire_api") or wire_api
                requires_openai_auth = bool(
                    provider.get("requires_openai_auth", requires_openai_auth)
                )
        if not provider_name and base_url:
            for candidate_name, candidate in providers.items():
                if optional_text(candidate, "base_url") == base_url:
                    provider_name = str(candidate_name)
                    wire_api = optional_text(candidate, "wire_api") or wire_api
                    requires_openai_auth = bool(
                        candidate.get("requires_openai_auth", requires_openai_auth)
                    )
                    break

    if not base_url:
        raise AgswitchError(
            f"Could not determine Codex base URL from {auth_path} and {config_path}"
        )
    if not provider_name:
        provider_name = default_provider_name or slug_from_url(base_url)

    return CodexProfile(
        api_key=api_key,
        base_url=base_url,
        provider_name=provider_name,
        wire_api=wire_api,
        requires_openai_auth=requires_openai_auth,
        auth_mode=auth_mode,
    )


def resolve_import_source(
    agent: str, source: Optional[Path], app_paths: AppPaths
) -> Profile:
    base = source.resolve() if source else app_paths.home
    name = "imported"

    if agent == "claude":
        return Profile(name=name, claude=parse_claude_settings(resolve_claude_settings(base)))
    if agent == "codex":
        auth_path, config_path = resolve_codex_paths(base)
        return Profile(name=name, codex=parse_codex_files(auth_path, config_path))
    if agent == "all":
        profile = Profile(name=name)
        profile.claude = parse_claude_settings(resolve_claude_settings(base))
        auth_path, config_path = resolve_codex_paths(base)
        profile.codex = parse_codex_files(auth_path, config_path)
        return profile
    raise AgswitchError(f"Unsupported agent: {agent}")


def resolve_claude_settings(source: Path) -> Path:
    candidates = []
    if source.is_file():
        candidates.append(source)
    else:
        candidates.extend([source / "settings.json", source / ".claude" / "settings.json"])
    for candidate in candidates:
        if candidate.exists():
            return candidate
    raise AgswitchError(f"Could not find Claude settings.json from {source}")


def resolve_codex_paths(source: Path) -> Tuple[Path, Path]:
    if source.is_dir():
        direct_auth = source / "auth.json"
        direct_config = source / "config.toml"
        nested_auth = source / ".codex" / "auth.json"
        nested_config = source / ".codex" / "config.toml"
        if direct_auth.exists() or direct_config.exists():
            return direct_auth, direct_config
        if nested_auth.exists() or nested_config.exists():
            return nested_auth, nested_config
        raise AgswitchError(f"Could not find Codex auth/config under {source}")

    if source.suffix == ".toml":
        config_path = source
        auth_path = source.with_name("auth.json")
        return auth_path, config_path

    auth_path = source
    config_path = source.with_name("config.toml")
    return auth_path, config_path


def import_profile(
    app_paths: AppPaths,
    name: str,
    agent: str,
    *,
    source: Optional[Path],
    overwrite: bool,
) -> Profile:
    imported = resolve_import_source(agent, source, app_paths)
    imported.name = name
    return add_profile(app_paths, imported, overwrite=overwrite)


def file_profile_name(path: Path, prefix: str) -> Optional[str]:
    filename = path.name
    if not filename.startswith(prefix) or not filename.endswith(".json"):
        return None
    name = filename[len(prefix) : -5].strip()
    return name or None


def build_profile_prefix(basename: str, separator: str) -> str:
    return f"{basename}{separator}"


def discover_existing_profiles(
    app_paths: AppPaths,
    *,
    include_current: bool = False,
    separator: str = DEFAULT_IMPORT_SEPARATOR,
) -> Dict[str, Profile]:
    discovered: Dict[str, Profile] = {}
    claude_prefix = build_profile_prefix(CLAUDE_PROFILE_BASENAME, separator)
    codex_prefix = build_profile_prefix(CODEX_PROFILE_BASENAME, separator)

    claude_dir = app_paths.claude_settings.parent
    if claude_dir.exists():
        for path in sorted(claude_dir.glob(f"{claude_prefix}*.json")):
            name = file_profile_name(path, claude_prefix)
            if not name:
                continue
            profile = discovered.setdefault(name, Profile(name=name))
            if profile.claude is not None:
                raise AgswitchError(f"Duplicate Claude profile name discovered: {name}")
            profile.claude = parse_claude_settings(path)

    codex_dir = app_paths.codex_auth.parent
    if codex_dir.exists():
        for path in sorted(codex_dir.glob(f"{codex_prefix}*.json")):
            name = file_profile_name(path, codex_prefix)
            if not name:
                continue
            profile = discovered.setdefault(name, Profile(name=name))
            if profile.codex is not None:
                raise AgswitchError(f"Duplicate Codex profile name discovered: {name}")
            profile.codex = parse_codex_files(
                path,
                app_paths.codex_config,
                default_provider_name=name,
            )

    if include_current:
        current = discovered.setdefault("current", Profile(name="current"))
        if current.claude is None and app_paths.claude_settings.exists():
            current.claude = parse_claude_settings(app_paths.claude_settings)
        if current.codex is None and app_paths.codex_auth.exists():
            current.codex = parse_codex_files(app_paths.codex_auth, app_paths.codex_config)

    return discovered


def import_existing_profiles(
    app_paths: AppPaths,
    *,
    overwrite: bool,
    include_current: bool = False,
    separator: str = DEFAULT_IMPORT_SEPARATOR,
) -> List[Profile]:
    discovered = discover_existing_profiles(
        app_paths,
        include_current=include_current,
        separator=separator,
    )
    if not discovered:
        raise AgswitchError("No existing named profile files were found.")

    profiles = load_profiles(app_paths)
    imported: List[Profile] = []
    for name in sorted(discovered):
        merged = merge_profile(
            profiles.get(name),
            discovered[name],
            overwrite=overwrite,
        )
        profiles[name] = merged
        imported.append(merged)
    save_profiles(app_paths, profiles)
    return imported


def profile_matches_claude(live: ClaudeProfile, candidate: ClaudeProfile) -> bool:
    return (
        live.auth_token == candidate.auth_token
        and live.base_url == candidate.base_url
    )


def profile_matches_codex(live: CodexProfile, candidate: CodexProfile) -> bool:
    return (
        live.api_key == candidate.api_key
        and live.base_url == candidate.base_url
        and live.provider_name == candidate.provider_name
    )


def find_matching_profile(
    profiles: Dict[str, Profile],
    *,
    claude: Optional[ClaudeProfile] = None,
    codex: Optional[CodexProfile] = None,
) -> Optional[str]:
    for name, profile in sorted(profiles.items()):
        if claude and profile.claude and profile_matches_claude(claude, profile.claude):
            return name
        if codex and profile.codex and profile_matches_codex(codex, profile.codex):
            return name
    return None


def get_status(app_paths: AppPaths, agent: str) -> List[LiveStatus]:
    profiles = load_profiles(app_paths)
    results: List[LiveStatus] = []
    if agent in ("claude", "all"):
        live = parse_claude_settings(app_paths.claude_settings)
        results.append(
            LiveStatus(
                agent="claude",
                profile_name=find_matching_profile(profiles, claude=live),
                details={
                    "base_url": live.base_url,
                    "token": mask_secret(live.auth_token),
                    "model": live.model or "-",
                },
                paths=[app_paths.claude_settings],
            )
        )
    if agent in ("codex", "all"):
        live = parse_codex_files(app_paths.codex_auth, app_paths.codex_config)
        results.append(
            LiveStatus(
                agent="codex",
                profile_name=find_matching_profile(profiles, codex=live),
                details={
                    "base_url": live.base_url,
                    "api_key": mask_secret(live.api_key),
                    "provider": live.provider_name,
                    "wire_api": live.wire_api,
                },
                paths=[app_paths.codex_auth, app_paths.codex_config],
            )
        )
    return results


def apply_profile(
    app_paths: AppPaths,
    agent: str,
    profile_name: str,
    *,
    dry_run: bool = False,
) -> Dict[str, List[Path]]:
    profiles = load_profiles(app_paths)
    if profile_name not in profiles:
        raise AgswitchError(f"Unknown profile: {profile_name}")
    profile = profiles[profile_name]

    backups: Dict[str, List[Path]] = {"claude": [], "codex": []}
    if agent in ("claude", "all"):
        if not profile.claude:
            raise AgswitchError(f"Profile '{profile_name}' has no claude section")
        backups["claude"] = apply_claude_profile(
            app_paths, profile.claude, dry_run=dry_run
        )
    if agent in ("codex", "all"):
        if not profile.codex:
            raise AgswitchError(f"Profile '{profile_name}' has no codex section")
        backups["codex"] = apply_codex_profile(
            app_paths, profile.codex, dry_run=dry_run
        )
    return backups


def apply_claude_profile(
    app_paths: AppPaths,
    profile: ClaudeProfile,
    *,
    dry_run: bool,
) -> List[Path]:
    path = app_paths.claude_settings
    data = read_json_file(path, missing_ok=True)
    env = data.get("env")
    if not isinstance(env, dict):
        env = {}
        data["env"] = env

    env["ANTHROPIC_AUTH_TOKEN"] = profile.auth_token
    env["ANTHROPIC_BASE_URL"] = profile.base_url
    if profile.disable_nonessential_traffic is not None:
        env["CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC"] = (
            profile.disable_nonessential_traffic
        )
    if profile.model:
        data["model"] = profile.model

    backup = backup_file(path, app_paths.backup_root)
    if not dry_run:
        atomic_write_json(path, data)
    return [item for item in [backup] if item]


def apply_codex_profile(
    app_paths: AppPaths,
    profile: CodexProfile,
    *,
    dry_run: bool,
) -> List[Path]:
    auth_data = read_json_file(app_paths.codex_auth, missing_ok=True)
    auth_data["auth_mode"] = profile.auth_mode
    auth_data["OPENAI_API_KEY"] = profile.api_key

    config_doc = tomlkit.document()
    if app_paths.codex_config.exists():
        try:
            config_doc = tomlkit.parse(app_paths.codex_config.read_text(encoding="utf-8"))
        except Exception as exc:
            raise AgswitchError(
                f"Invalid TOML in {app_paths.codex_config}: {exc}"
            ) from exc

    config_doc["model_provider"] = profile.provider_name
    providers = config_doc.get("model_providers")
    if providers is None:
        providers = tomlkit.table()
        config_doc["model_providers"] = providers
    provider = providers.get(profile.provider_name)
    if provider is None:
        provider = tomlkit.table()
        providers[profile.provider_name] = provider
    provider["name"] = profile.provider_name
    provider["base_url"] = profile.base_url
    provider["requires_openai_auth"] = profile.requires_openai_auth
    provider["wire_api"] = profile.wire_api

    auth_backup = backup_file(app_paths.codex_auth, app_paths.backup_root)
    config_backup = backup_file(app_paths.codex_config, app_paths.backup_root)
    if not dry_run:
        atomic_write_json(app_paths.codex_auth, auth_data)
        atomic_write_text(app_paths.codex_config, tomlkit.dumps(config_doc))
    return [item for item in [auth_backup, config_backup] if item]
