"""
auth.py — Login flow supporting OneFirewall AI and any OpenAI-compatible provider.

Provider flow:
  1) OneFirewall AI  → hardcoded endpoint, ask API key, list models (workspaces)
  2) Custom provider → ask base URL, ask API key, ask model name; test then save
"""

import getpass
import sys

import requests
from rich.console import Console
from rich.prompt import Prompt
from rich.panel import Panel
from rich.rule import Rule

from onefirewall_ai.config import load_access, save_access

console = Console()

ONEFIREWALL_ENDPOINT = "https://onefirewall.ai"


# ─────────────────────────────────────────────────────────────────────────────
# Shared helpers
# ─────────────────────────────────────────────────────────────────────────────

def _get_available_models(endpoint: str, api_key: str) -> list[str]:
    """
    Fetch model IDs from  <endpoint>/api/v1/models  (OneFirewall gateway format)
    or  <endpoint>/v1/models  (standard OpenAI format).

    Returns a list of model ID strings, or [] on failure.
    """
    headers = {"Authorization": f"Bearer {api_key}"}

    # Try OneFirewall path first, then OpenAI path
    for path in ("/api/v1/models", "/v1/models"):
        url = endpoint.rstrip("/") + path
        try:
            resp = requests.get(url, headers=headers, timeout=10)
            if resp.status_code == 200:
                data = resp.json().get("data", [])
                ids = [m["id"] for m in data if "id" in m]
                if ids:
                    return ids
        except Exception:
            continue
    return []


def _test_connection(endpoint: str, api_key: str, model: str) -> tuple[bool, str]:
    """
    Send a minimal chat completion to verify the endpoint + key + model combo.
    Tries   <endpoint>/api/v1/chat/completions   then   <endpoint>/v1/chat/completions.
    Returns (success: bool, message: str).
    """
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": "ping"}],
        "stream": False,
    }

    for path in ("/api/v1/chat/completions", "/v1/chat/completions"):
        url = endpoint.rstrip("/") + path
        try:
            resp = requests.post(url, headers=headers, json=payload, timeout=15)
            if resp.status_code == 200:
                return True, "Connection successful ✓"
            elif resp.status_code == 401:
                return False, "Invalid API key (401 Unauthorized)"
            elif resp.status_code == 403:
                return False, "Access forbidden (403) — check your API key permissions"
            elif resp.status_code == 404:
                # This path doesn't exist, try the next one
                continue
            else:
                return False, f"HTTP {resp.status_code} — {resp.text[:200]}"
        except requests.exceptions.ConnectionError:
            return False, f"Cannot connect to '{endpoint}' — is the server running?"
        except requests.exceptions.Timeout:
            return False, "Connection timed out after 15 seconds"
        except Exception as exc:
            return False, f"Unexpected error: {exc}"

    return False, f"No valid API path found on '{endpoint}'"


def _prompt_api_key(label: str = "API Key (hidden)") -> str:
    """Prompt for a hidden API key; exit gracefully on Ctrl-C."""
    try:
        key = getpass.getpass(f"{label}: ").strip()
    except (KeyboardInterrupt, EOFError):
        console.print("\n[yellow]Login cancelled.[/yellow]")
        sys.exit(0)
    if not key:
        console.print("[red]✗ API key cannot be empty.[/red]")
        sys.exit(1)
    return key


def _select_model_from_list(models: list[str], current: str) -> str:
    """
    Show a numbered model list and return the user's selection.
    Falls back to `current` if the list has only one entry (auto-selects).
    """
    if len(models) == 1:
        console.print(
            f"[dim]Only one model available — auto-selecting "
            f"[bold cyan]{models[0]}[/bold cyan][/dim]"
        )
        return models[0]

    console.print()
    console.print(f"[bold cyan]Found {len(models)} available models:[/bold cyan]")
    for i, m in enumerate(models, 1):
        marker = " [dim]← current[/dim]" if m == current else ""
        console.print(f"  [bold white]{i})[/bold white] {m}{marker}")

    default_idx = "1"
    if current in models:
        default_idx = str(models.index(current) + 1)

    choice = Prompt.ask(
        "\n[bold]Select model index[/bold]",
        choices=[str(i) for i in range(1, len(models) + 1)],
        default=default_idx,
        console=console,
    )
    selected = models[int(choice) - 1]
    console.print(f"[dim]Selected: [bold cyan]{selected}[/bold cyan][/dim]")
    return selected


# ─────────────────────────────────────────────────────────────────────────────
# Provider-specific flows
# ─────────────────────────────────────────────────────────────────────────────

def _login_onefirewall(existing: dict) -> dict:
    """
    OneFirewall AI login:
      - Endpoint is fixed (https://onefirewall.ai)
      - Ask API key
      - List workspaces / models → user selects
      - Test connection → save
    """
    console.print()
    console.print(Rule("[bold cyan]OneFirewall AI[/bold cyan]", style="cyan"))
    console.print(
        "[dim]Connecting to the OneFirewall AI managed gateway.[/dim]\n"
        "[dim]Your API key controls which workspaces you can access.[/dim]"
    )
    console.print()

    endpoint = ONEFIREWALL_ENDPOINT

    # ── API Key ───────────────────────────────────────────────────────────────
    console.print("[dim]Paste your OneFirewall AI API key below.[/dim]")
    api_key = _prompt_api_key("OneFirewall API Key (hidden)")

    # ── Fetch models (workspaces) ─────────────────────────────────────────────
    console.print()
    with console.status("[bold cyan]Fetching available workspaces / models…[/bold cyan]", spinner="dots"):
        models = _get_available_models(endpoint, api_key)

    current_model = existing.get("model") or ""

    if models:
        selected_model = _select_model_from_list(models, current_model)
    else:
        console.print(
            "[yellow]! Could not fetch models from OneFirewall AI.[/yellow]\n"
            "[dim]Tip: Verify your API key; the gateway may be temporarily unavailable.[/dim]"
        )
        # Fallback: keep existing or ask
        fallback = current_model or "openai/gpt-4o"
        selected_model = Prompt.ask(
            "[bold]Enter model name manually[/bold]",
            default=fallback,
            console=console,
        ).strip()

    # ── Test connection ───────────────────────────────────────────────────────
    console.print()
    with console.status(
        f"[bold cyan]Testing connection with '{selected_model}'…[/bold cyan]", spinner="dots"
    ):
        ok, msg = _test_connection(endpoint, api_key, model=selected_model)

    if ok:
        console.print(f"[bold green]✓ {msg}[/bold green]")
    else:
        console.print(f"[bold red]✗ {msg}[/bold red]")
        console.print("[dim]Check your API key and try again.[/dim]")
        sys.exit(1)

    return {
        **existing,
        "provider": "onefirewall",
        "endpoint": endpoint,
        "api_key": api_key,
        "model": selected_model,
    }


def _login_custom_provider(existing: dict) -> dict:
    """
    Generic OpenAI-compatible provider login:
      - Ask provider display name (label only)
      - Ask API base URL
      - Ask API key
      - Ask model name (free text)
      - Test connection → save
    """
    console.print()
    console.print(Rule("[bold magenta]Custom OpenAI-Compatible Provider[/bold magenta]", style="magenta"))
    console.print(
        "[dim]Configure any provider that exposes an OpenAI-compatible API.[/dim]\n"
        "[dim]Examples: OpenAI, Groq, Together AI, Ollama, Azure OpenAI, LM Studio…[/dim]"
    )
    console.print()

    # ── Provider label ─────────────────────────────────────────────────────────
    provider_label = Prompt.ask(
        "[bold]Provider name[/bold] [dim](label, e.g. OpenAI / Groq / Ollama)[/dim]",
        default=existing.get("provider_label") or "Custom",
        console=console,
    ).strip()

    # ── Base URL ───────────────────────────────────────────────────────────────
    console.print()
    console.print(
        "[dim]Enter the base URL of the provider's API.[/dim]\n"
        "[dim]Examples:[/dim]\n"
        "  [dim]• OpenAI       →[/dim] [cyan]https://api.openai.com[/cyan]\n"
        "  [dim]• Groq         →[/dim] [cyan]https://api.groq.com/openai[/cyan]\n"
        "  [dim]• Ollama local →[/dim] [cyan]http://localhost:11434[/cyan]"
    )

    current_ep = existing.get("endpoint") or ""
    # Don't reuse the onefirewall endpoint as default for custom
    if current_ep == ONEFIREWALL_ENDPOINT:
        current_ep = ""

    endpoint = Prompt.ask(
        "\n[bold]API Base URL[/bold]",
        default=current_ep or "https://api.openai.com",
        console=console,
    ).strip().rstrip("/")

    if endpoint and not endpoint.startswith(("http://", "https://")):
        endpoint = "https://" + endpoint

    # ── API Key ───────────────────────────────────────────────────────────────
    console.print()
    console.print("[dim]Get your API key from the provider's dashboard.[/dim]")
    api_key = _prompt_api_key(f"{provider_label} API Key (hidden)")

    # ── Model name ────────────────────────────────────────────────────────────
    console.print()

    # Try to fetch models automatically — many providers support /v1/models
    with console.status("[bold cyan]Looking for available models…[/bold cyan]", spinner="dots"):
        auto_models = _get_available_models(endpoint, api_key)

    current_model = existing.get("model") or ""
    if auto_models:
        selected_model = _select_model_from_list(auto_models, current_model)
    else:
        console.print("[dim]Could not auto-fetch model list — please enter the model name manually.[/dim]")
        console.print(
            "[dim]Examples: gpt-4o, gpt-3.5-turbo, llama3-8b-8192, mixtral-8x7b-32768[/dim]"
        )
        selected_model = Prompt.ask(
            "[bold]Model name[/bold]",
            default=current_model or "gpt-4o",
            console=console,
        ).strip()

    # ── Test connection ────────────────────────────────────────────────────────
    console.print()
    with console.status(
        f"[bold cyan]Testing connection to {provider_label} with '{selected_model}'…[/bold cyan]",
        spinner="dots",
    ):
        ok, msg = _test_connection(endpoint, api_key, model=selected_model)

    if ok:
        console.print(f"[bold green]✓ {msg}[/bold green]")
    else:
        console.print(f"[bold red]✗ {msg}[/bold red]")
        console.print("[dim]Double-check your base URL, API key, and model name, then try again.[/dim]")
        sys.exit(1)

    return {
        **existing,
        "provider": "custom",
        "provider_label": provider_label,
        "endpoint": endpoint,
        "api_key": api_key,
        "model": selected_model,
    }


# ─────────────────────────────────────────────────────────────────────────────
# Public entry point
# ─────────────────────────────────────────────────────────────────────────────

def login() -> None:
    """Interactive login — choose provider, configure, test, save."""
    console.print()
    console.print(Panel(
        "[bold cyan]OneFirewall AI — Configure Provider[/bold cyan]\n"
        "[dim]Connect to OneFirewall AI or any OpenAI-compatible API.[/dim]",
        border_style="cyan",
    ))
    console.print()

    existing = load_access()

    # ── Step 1: Choose provider ────────────────────────────────────────────────
    console.print("[bold]Select a provider:[/bold]\n")
    console.print("  [bold white]1)[/bold white] OneFirewall AI  [dim](managed AI gateway)[/dim]")
    console.print("  [bold white]2)[/bold white] Other           [dim](any OpenAI-compatible provider)[/dim]")
    console.print()

    # Pre-select based on existing config
    current_provider = existing.get("provider", "onefirewall")
    default_choice = "1" if current_provider == "onefirewall" else "2"

    choice = Prompt.ask(
        "[bold]Provider[/bold]",
        choices=["1", "2"],
        default=default_choice,
        console=console,
    )

    # ── Step 2: Provider-specific flow ────────────────────────────────────────
    if choice == "1":
        new_config = _login_onefirewall(existing)
    else:
        new_config = _login_custom_provider(existing)

    # ── Step 3: Save ──────────────────────────────────────────────────────────
    save_access(new_config)

    # Memory sync is a OneFirewall AI gateway feature — skip for custom providers
    if new_config.get("provider") == "onefirewall":
        from onefirewall_ai.utils import sync_memory
        sync_memory()

    provider_display = (
        "OneFirewall AI"
        if new_config.get("provider") == "onefirewall"
        else new_config.get("provider_label", "Custom provider")
    )

    console.print()
    console.print(Panel(
        f"[bold green]✓ Logged in successfully![/bold green]\n"
        f"[dim]Provider:[/dim]  [cyan]{provider_display}[/cyan]\n"
        f"[dim]Endpoint:[/dim]  [cyan]{new_config['endpoint']}[/cyan]\n"
        f"[dim]Model:[/dim]     [cyan]{new_config['model']}[/cyan]\n"
        f"[dim]Config:[/dim]    [cyan]~/.onefirewall-ai/access.json[/cyan]",
        border_style="green",
    ))
    console.print()
    console.print("[dim]Run [bold]onefirewall-ai[/bold] to start chatting.[/dim]")
    console.print()
