import os
import logging
import asyncio
from typing import AsyncIterator, Dict, Any, Optional
from backend.app.llm.base import BaseLLMProvider
from backend.app.llm.local import LocalLLMProvider
from backend.app.config import settings

logger = logging.getLogger(__name__)

MODEL_REPO = "bartowski/SmolLM2-360M-Instruct-GGUF"
MODEL_FILE = "SmolLM2-360M-Instruct-Q4_K_M.gguf"


class LLMManager:
    def __init__(self):
        self.provider: Optional[BaseLLMProvider] = None

    async def ensure_model_downloaded(self) -> bool:
        model_path = os.path.abspath(settings.LOCAL_MODEL_PATH)
        if os.path.exists(model_path):
            logger.info(f"Model found at {model_path}")
            return True

        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        logger.info(f"Downloading {MODEL_FILE} (~180MB) from {MODEL_REPO}...")
        try:
            from huggingface_hub import hf_hub_download

            def download():
                return hf_hub_download(
                    repo_id=MODEL_REPO, filename=MODEL_FILE,
                    local_dir=os.path.dirname(model_path),
                    local_dir_use_symlinks=False,
                )

            await asyncio.to_thread(download)
            logger.info("Model download complete!")
            return True
        except Exception as e:
            logger.error(f"Model download failed: {e}")
            return False

    async def setup_provider(self):
        logger.info("Initializing local provider...")
        await self.ensure_model_downloaded()
        local = LocalLLMProvider()
        await local.initialize()
        self.provider = local
        return local

    async def get_active_provider(self) -> BaseLLMProvider:
        if not self.provider:
            await self.setup_provider()
        return self.provider

    async def generate(self, *args, **kwargs) -> Dict[str, Any]:
        p = await self.get_active_provider()
        return await p.generate(*args, **kwargs)

    async def generate_stream(self, *args, **kwargs) -> AsyncIterator[str]:
        p = await self.get_active_provider()
        return await p.generate_stream(*args, **kwargs)

    async def get_status_info(self) -> Dict[str, Any]:
        if self.provider:
            return self.provider.get_info()
        return {
            "provider": "none",
            "initialized": False,
            "model_path": os.path.abspath(settings.LOCAL_MODEL_PATH),
            "context_length": settings.DEFAULT_CONTEXT_LENGTH,
        }


llm_manager = LLMManager()