"""Cloud-first hybrid startup -- value starts IMMEDIATELY.

The :class:`HybridStartup` orchestrates a two-phase initialisation:

1. **Fast phase** (<1 s): check cloud API keys, make them available to the
   router immediately so the proxy can serve requests.
2. **Background phase** (async): discover Ollama, probe local models, and
   progressively register them with the router so that traffic automatically
   shifts to local backends as they become available.

There is never a "waiting for setup" state.  If cloud keys exist the proxy is
ready instantly; local discovery enriches it in parallel.
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import time
from typing import TYPE_CHECKING

from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from llmhosts.discovery.ollama import OllamaDiscovery
    from llmhosts.keys.manager import KeyManager
    from llmhosts.router.engine import Router

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Status model
# ---------------------------------------------------------------------------


class HybridStatus(BaseModel):
    """Current status of local vs cloud availability."""

    cloud_ready: bool = False
    cloud_providers: list[str] = Field(default_factory=list)
    local_ready: bool = False
    local_models: list[str] = Field(default_factory=list)
    mode: str = "none"  # "cloud-only", "local-only", "hybrid", "none"
    recommendations: list[str] = Field(default_factory=list)


# ---------------------------------------------------------------------------
# Hybrid startup orchestrator
# ---------------------------------------------------------------------------


class HybridStartup:
    """Implements the cloud-first hybrid onboarding experience.

    Value starts IMMEDIATELY with cloud API keys.
    Local Ollama discovered in parallel.
    Traffic automatically shifts to local as models become available.
    Never a "waiting for setup" state.

    Usage::

        startup = HybridStartup(keys=key_mgr, ollama=ollama, router=router)
        status = await startup.initialize()
        # proxy is ready NOW -- background discovery continues
        await startup.background_discovery()
    """

    def __init__(
        self,
        keys: KeyManager,
        ollama: OllamaDiscovery,
        router: Router,
        ollama_host: str = "http://127.0.0.1:11434",
    ) -> None:
        self._keys = keys
        self._ollama = ollama
        self._router = router
        self._ollama_host = ollama_host
        self._local_ready: bool = False
        self._cloud_ready: bool = False
        self._cloud_providers: list[str] = []
        self._local_models: list[str] = []
        self._discovery_task: asyncio.Task[None] | None = None
        self._shutdown: bool = False

    # ------------------------------------------------------------------
    # Initialisation (fast path)
    # ------------------------------------------------------------------

    async def initialize(self) -> HybridStatus:
        """Concurrent initialisation.

        - Check cloud keys (fast, <1 s).
        - Start Ollama discovery (may take a few seconds).
        - Return status immediately with what's available.

        Returns
        -------
        HybridStatus
            Snapshot of availability right after the fast phase.
        """
        start = time.monotonic()

        # Run cloud key check and a quick Ollama probe concurrently
        cloud_task = asyncio.create_task(self._check_cloud_keys())
        local_task = asyncio.create_task(self._quick_local_probe())

        await asyncio.gather(cloud_task, local_task, return_exceptions=True)

        elapsed = (time.monotonic() - start) * 1000
        status = self.get_status()

        logger.info(
            "HybridStartup initialised in %.0fms: mode=%s cloud=%s local=%d model(s)",
            elapsed,
            status.mode,
            status.cloud_providers,
            len(status.local_models),
        )

        return status

    # ------------------------------------------------------------------
    # Background discovery
    # ------------------------------------------------------------------

    async def background_discovery(self) -> None:
        """Background task that continues discovering local capabilities.

        When new local models are found, updates the router automatically.
        Designed to run as a long-lived ``asyncio.Task``.
        """
        # Initial deep discovery
        await self._deep_local_discovery()

        # Then poll periodically for changes (new models pulled, Ollama restarted, etc.)
        poll_interval = 30.0  # seconds between re-polls
        while not self._shutdown:
            try:
                await asyncio.sleep(poll_interval)
                if self._shutdown:
                    break
                await self._deep_local_discovery()
                # Exponential back-off on success (up to 5 minutes)
                poll_interval = min(poll_interval * 1.5, 300.0)
            except asyncio.CancelledError:
                break
            except Exception as exc:
                logger.debug("Background discovery error: %s", exc)
                # Back-off on error
                poll_interval = min(poll_interval * 2, 300.0)

    def start_background_discovery(self) -> asyncio.Task[None]:
        """Launch :meth:`background_discovery` as an ``asyncio.Task``.

        Returns the task so the caller can await or cancel it.
        """
        self._discovery_task = asyncio.create_task(
            self.background_discovery(),
            name="llmhosts-hybrid-discovery",
        )
        return self._discovery_task

    async def shutdown(self) -> None:
        """Signal background discovery to stop and await its completion."""
        self._shutdown = True
        if self._discovery_task is not None and not self._discovery_task.done():
            self._discovery_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._discovery_task
        logger.debug("HybridStartup shutdown complete")

    # ------------------------------------------------------------------
    # Status
    # ------------------------------------------------------------------

    def get_status(self) -> HybridStatus:
        """Current status of local vs cloud availability."""
        recommendations: list[str] = []

        if not self._cloud_ready and not self._local_ready:
            mode = "none"
            recommendations.append("Add a cloud API key: llmhosts keys add openai <key>")
            recommendations.append("Or install Ollama: https://ollama.com")
        elif self._cloud_ready and self._local_ready:
            mode = "hybrid"
            recommendations.append("Hybrid mode active -- requests route to cheapest option")
        elif self._cloud_ready:
            mode = "cloud-only"
            recommendations.append("Install Ollama for free local inference: https://ollama.com")
            if not self._local_models:
                recommendations.append("Pull a model: ollama pull llama3.2")
        else:
            mode = "local-only"
            recommendations.append("Add a cloud key for fallback: llmhosts keys add openai <key>")

        return HybridStatus(
            cloud_ready=self._cloud_ready,
            cloud_providers=list(self._cloud_providers),
            local_ready=self._local_ready,
            local_models=list(self._local_models),
            mode=mode,
            recommendations=recommendations,
        )

    # ------------------------------------------------------------------
    # Internal: cloud keys
    # ------------------------------------------------------------------

    async def _check_cloud_keys(self) -> None:
        """Check which cloud API keys are configured (fast, sync operation)."""
        try:
            providers = self._keys.list_providers()
            self._cloud_providers = [p.provider for p in providers]
            self._cloud_ready = len(self._cloud_providers) > 0

            if self._cloud_ready:
                # Register cloud backends with the router
                self._register_cloud_backends()
                logger.info("Cloud keys ready: %s", self._cloud_providers)
            else:
                logger.info("No cloud API keys configured")
        except Exception as exc:
            logger.warning("Cloud key check failed: %s", exc)
            self._cloud_ready = False

    def _register_cloud_backends(self) -> None:
        """Register cloud providers as backends in the router."""
        from llmhosts.router.models import BackendInfo

        # Provider -> API base URL mapping
        provider_urls: dict[str, str] = {
            "openai": "https://api.openai.com/v1",
            "anthropic": "https://api.anthropic.com/v1",
            "openrouter": "https://openrouter.ai/api/v1",
            "google": "https://generativelanguage.googleapis.com/v1",
            "mistral": "https://api.mistral.ai/v1",
            "groq": "https://api.groq.com/openai/v1",
            "together": "https://api.together.xyz/v1",
            "deepseek": "https://api.deepseek.com/v1",
        }

        # Model catalogues (representative models per provider)
        provider_models: dict[str, list[str]] = {
            "openai": ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo", "gpt-3.5-turbo", "o1", "o1-mini", "o3-mini"],
            "anthropic": [
                "claude-3-5-sonnet-20241022",
                "claude-3-5-haiku-20241022",
                "claude-3-opus-20240229",
                "claude-3-haiku-20240307",
            ],
            "google": ["gemini-1.5-pro", "gemini-1.5-flash"],
            "mistral": ["mistral-large-latest", "mistral-small-latest"],
            "groq": ["llama-3.1-70b-versatile", "llama-3.1-8b-instant", "mixtral-8x7b-32768"],
            "together": ["meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo"],
            "deepseek": ["deepseek-chat", "deepseek-coder"],
            "openrouter": [],  # universal fallback, accepts any model
        }

        for provider in self._cloud_providers:
            url = provider_urls.get(provider, "")
            if not url:
                logger.debug("Unknown provider URL for '%s', skipping backend registration", provider)
                continue

            models = provider_models.get(provider, [])
            backend = BackendInfo(
                backend_type=provider,
                url=url,
                models=models,
                is_local=False,
                is_healthy=True,
            )
            self._router.register_backend(backend)

    # ------------------------------------------------------------------
    # Internal: local Ollama
    # ------------------------------------------------------------------

    async def _quick_local_probe(self) -> None:
        """Quick probe: is Ollama running?  If yes, list models."""
        try:
            available = await self._ollama.is_available()
            if available:
                models = await self._ollama.list_models()
                model_names = [m.name for m in models]
                self._local_models = model_names
                self._local_ready = len(model_names) > 0
                if self._local_ready:
                    self._register_ollama_backend(model_names)
                    logger.info("Ollama quick probe: %d model(s) available", len(model_names))
                else:
                    logger.info("Ollama running but no models installed")
            else:
                logger.info("Ollama not reachable at %s", self._ollama_host)
        except Exception as exc:
            logger.debug("Quick Ollama probe failed: %s", exc)

    async def _deep_local_discovery(self) -> None:
        """Deep discovery: re-check Ollama, update model list and router."""
        try:
            available = await self._ollama.is_available()
            if not available:
                if self._local_ready:
                    # Ollama went away
                    logger.warning("Ollama became unreachable -- removing local backend")
                    self._local_ready = False
                    self._local_models = []
                    self._router.unregister_backend("ollama")
                return

            models = await self._ollama.list_models()
            model_names = [m.name for m in models]
            prev_models = set(self._local_models)
            new_models = set(model_names)

            if new_models != prev_models:
                added = new_models - prev_models
                removed = prev_models - new_models
                if added:
                    logger.info("New local models discovered: %s", added)
                if removed:
                    logger.info("Local models removed: %s", removed)

                self._local_models = model_names
                self._local_ready = len(model_names) > 0
                if self._local_ready:
                    self._register_ollama_backend(model_names)
                else:
                    self._router.unregister_backend("ollama")

        except Exception as exc:
            logger.debug("Deep local discovery error: %s", exc)

    def _register_ollama_backend(self, models: list[str]) -> None:
        """Register or update the Ollama backend in the router."""
        from llmhosts.router.models import BackendInfo

        backend = BackendInfo(
            backend_type="ollama",
            url=self._ollama_host,
            models=models,
            is_local=True,
            is_healthy=True,
        )
        self._router.register_backend(backend)
