"""
AI Provider Service - Handles communication with different AI backends.
Supports: Anthropic (Claude), OpenAI, Ollama, Cloudflare Workers AI, and custom OpenAI-compatible endpoints.
"""

from typing import Literal, AsyncGenerator
import httpx
from pydantic import BaseModel

from app.core.config import settings

# Try importing provider SDKs
try:
    import anthropic
    HAS_ANTHROPIC = True
except ImportError:
    HAS_ANTHROPIC = False

try:
    import openai
    HAS_OPENAI = True
except ImportError:
    HAS_OPENAI = False


class ProviderConfig(BaseModel):
    """Provider configuration from frontend."""
    id: str
    name: str
    type: str  # Provider type (validated against model registry)
    enabled: bool
    endpoint: str | None = None
    api_key: str | None = None
    account_id: str | None = None  # For Cloudflare
    model: str


class Message(BaseModel):
    """Chat message with optional image support."""
    role: Literal["user", "assistant", "system"]
    content: str
    images: list[str] | None = None  # List of base64-encoded images


async def complete_with_anthropic(
    messages: list[Message],
    config: ProviderConfig,
    system_prompt: str | None = None,
) -> str:
    """Complete chat using Anthropic Claude API."""
    if not HAS_ANTHROPIC:
        raise RuntimeError("anthropic package not installed. Run: uv add anthropic")
    
    # Use env var as fallback
    api_key = config.api_key or settings.anthropic_api_key
    if not api_key:
        raise ValueError("Anthropic API key is required. Set ANTHROPIC_API_KEY env var or provide in settings.")
    
    client = anthropic.Anthropic(api_key=api_key)
    
    # Convert messages to Anthropic format
    anthropic_messages = []
    for msg in messages:
        if msg.role != "system":
            anthropic_messages.append({
                "role": msg.role,
                "content": msg.content,
            })
    
    # Build system prompt from system messages
    system_parts = [system_prompt] if system_prompt else []
    for msg in messages:
        if msg.role == "system":
            system_parts.append(msg.content)
    
    response = client.messages.create(
        model=config.model,
        max_tokens=4096,
        system="\n\n".join(system_parts) if system_parts else None,
        messages=anthropic_messages,
    )
    
    return response.content[0].text


async def complete_with_openai(
    messages: list[Message],
    config: ProviderConfig,
    system_prompt: str | None = None,
) -> str:
    """Complete chat using OpenAI API."""
    if not HAS_OPENAI:
        raise RuntimeError("openai package not installed. Run: uv add openai")
    
    # Use env var as fallback
    api_key = config.api_key or settings.openai_api_key
    if not api_key:
        raise ValueError("OpenAI API key is required. Set OPENAI_API_KEY env var or provide in settings.")
    
    client = openai.OpenAI(api_key=api_key)
    
    # Build messages list
    openai_messages = []
    if system_prompt:
        openai_messages.append({"role": "system", "content": system_prompt})
    
    for msg in messages:
        openai_messages.append({
            "role": msg.role,
            "content": msg.content,
        })
    
    response = client.chat.completions.create(
        model=config.model,
        messages=openai_messages,
        max_tokens=4096,
    )
    
    return response.choices[0].message.content


async def complete_with_ollama(
    messages: list[Message],
    config: ProviderConfig,
    system_prompt: str | None = None,
) -> str:
    """Complete chat using Ollama local API (non-streaming). Supports multimodal (images)."""
    # Always default to localhost if no endpoint provided
    endpoint = (config.endpoint or "http://localhost:11434").rstrip('/')
    
    # Build messages list with optional images
    ollama_messages = []
    if system_prompt:
        ollama_messages.append({"role": "system", "content": system_prompt})
    
    for msg in messages:
        message_dict: dict = {
            "role": msg.role,
            "content": msg.content,
        }
        # Add images if present (for multimodal models like qwen3-vl, llava, etc.)
        if msg.images:
            message_dict["images"] = msg.images
        ollama_messages.append(message_dict)
    
    async with httpx.AsyncClient(timeout=120.0) as client:
        response = await client.post(
            f"{endpoint}/api/chat",
            json={
                "model": config.model,
                "messages": ollama_messages,
                "stream": False,
            },
        )
        response.raise_for_status()
        data = response.json()
        return data["message"]["content"]


async def stream_with_ollama(
    messages: list[Message],
    config: ProviderConfig,
    system_prompt: str | None = None,
) -> AsyncGenerator[str, None]:
    """Stream chat using Ollama local API. Supports multimodal (images).
    
    Handles Qwen3-style models that output 'thinking' tokens separately from 'content'.
    Wraps thinking in <think></think> tags so the frontend can render them.
    """
    import json
    
    endpoint = (config.endpoint or "http://localhost:11434").rstrip('/')
    
    # Build messages list with optional images
    ollama_messages = []
    if system_prompt:
        ollama_messages.append({"role": "system", "content": system_prompt})
    
    for msg in messages:
        message_dict: dict = {
            "role": msg.role,
            "content": msg.content,
        }
        # Add images if present (for multimodal models like qwen3-vl, llava, etc.)
        if msg.images:
            message_dict["images"] = msg.images
        ollama_messages.append(message_dict)
    
    async with httpx.AsyncClient(timeout=120.0) as client:
        async with client.stream(
            "POST",
            f"{endpoint}/api/chat",
            json={
                "model": config.model,
                "messages": ollama_messages,
                "stream": True,
            },
        ) as response:
            response.raise_for_status()
            
            # Ollama streams JSON lines (not SSE). Each line looks like:
            # {"model":"...","message":{"role":"assistant","content":"token"},"done":false}
            # Final line has {"done": true}
            async for line in response.aiter_lines():
                if not line:
                    continue
                try:
                    data = json.loads(line)
                except json.JSONDecodeError:
                    continue

                # Some implementations may send thinking tokens separately; if we
                # ever receive a dedicated field, wrap it in <think> tags. For now
                # we stream assistant content as-is and let the model emit tags.
                msg = data.get("message") or {}
                token = msg.get("content")
                if token:
                    yield token


async def stream_with_openai_compatible(
    messages: list[Message],
    config: ProviderConfig,
    system_prompt: str | None = None,
) -> AsyncGenerator[str, None]:
    """Stream chat using OpenAI-compatible API (OpenAI, XAI, Mistral, OpenRouter, etc)."""
    import json
    import logging
    
    logger = logging.getLogger(__name__)
    endpoint = config.endpoint.rstrip('/')
    
    # Build messages list
    api_messages = []
    if system_prompt:
        api_messages.append({"role": "system", "content": system_prompt})
    
    for msg in messages:
        api_messages.append({
            "role": msg.role,
            "content": msg.content,
        })
    
    headers = {
        "Authorization": f"Bearer {config.api_key[:8]}...",  # Mask API key in logs
        "Content-Type": "application/json",
    }
    
    # OpenRouter requires HTTP-Referer header
    if "openrouter" in endpoint:
        headers["HTTP-Referer"] = "https://nexus.local"
        headers["X-Title"] = "Nexus"
    
    # Restore full API key for actual request
    headers["Authorization"] = f"Bearer {config.api_key}"
    
    payload = {
        "model": config.model,
        "messages": api_messages,
        "stream": True,
    }
    
    logger.info(f"🚀 Streaming to {endpoint}/chat/completions with model={config.model}")
    
    async with httpx.AsyncClient(timeout=120.0) as client:
        async with client.stream(
            "POST",
            f"{endpoint}/chat/completions",
            headers=headers,
            json=payload,
        ) as response:
            # Check status before reading stream
            if response.status_code != 200:
                error_text = await response.aread()
                error_body = error_text.decode()
                logger.error(f"❌ HTTP {response.status_code}: {error_body[:500]}")
                
                # Try to extract user-friendly error message
                try:
                    error_json = json.loads(error_body)
                    if "error" in error_json:
                        if isinstance(error_json["error"], dict) and "message" in error_json["error"]:
                            user_message = error_json["error"]["message"]
                        elif isinstance(error_json["error"], str):
                            user_message = error_json["error"]
                        else:
                            user_message = str(error_json["error"])
                    else:
                        user_message = f"HTTP {response.status_code}: {error_body[:200]}"
                except:
                    user_message = f"HTTP {response.status_code}: {error_body[:200]}"
                
                raise httpx.HTTPStatusError(
                    user_message,
                    request=response.request,
                    response=response
                )
            
            async for line in response.aiter_lines():
                if not line or not line.startswith("data: "):
                    continue
                    
                data_str = line[6:]  # Remove "data: " prefix
                
                if data_str == "[DONE]":
                    break
                
                try:
                    data = json.loads(data_str)
                    choices = data.get("choices", [])
                    if not choices:
                        continue
                    
                    delta = choices[0].get("delta", {})
                    content = delta.get("content", "")
                    
                    if content:
                        yield content
                        
                except json.JSONDecodeError:
                    continue


async def complete_with_cloudflare(
    messages: list[Message],
    config: ProviderConfig,
    system_prompt: str | None = None,
) -> str:
    """Complete chat using Cloudflare Workers AI API."""
    # Use env vars as fallback
    api_key = config.api_key or settings.cloudflare_api_token
    account_id = config.account_id or settings.cloudflare_account_id
    
    if not api_key:
        raise ValueError("Cloudflare API token is required. Set CLOUDFLARE_API_TOKEN env var or provide in settings.")
    
    if not account_id:
        raise ValueError("Cloudflare Account ID is required. Set CLOUDFLARE_ACCOUNT_ID env var or provide in settings.")
    
    # Build messages list
    cf_messages = []
    if system_prompt:
        cf_messages.append({"role": "system", "content": system_prompt})
    
    for msg in messages:
        cf_messages.append({
            "role": msg.role,
            "content": msg.content,
        })
    
    # Cloudflare Workers AI endpoint
    # Model format: @cf/meta/llama-3.1-8b-instruct
    model_path = config.model.lstrip('@cf/')
    url = f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/@cf/{model_path}"
    
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    
    async with httpx.AsyncClient(timeout=120.0) as client:
        response = await client.post(
            url,
            headers=headers,
            json={
                "messages": cf_messages,
            },
        )
        response.raise_for_status()
        data = response.json()
        
        # Cloudflare returns { "result": { "response": "..." } }
        if "result" in data:
            return data["result"].get("response", str(data["result"]))
        return str(data)


async def stream_with_cloudflare(
    messages: list[Message],
    config: ProviderConfig,
    system_prompt: str | None = None,
) -> AsyncGenerator[str, None]:
    """Stream chat using Cloudflare Workers AI API with SSE."""
    import json
    
    # Use env vars as fallback
    api_key = config.api_key or settings.cloudflare_api_token
    account_id = config.account_id or settings.cloudflare_account_id
    
    if not api_key:
        raise ValueError("Cloudflare API token is required.")
    
    if not account_id:
        raise ValueError("Cloudflare Account ID is required.")
    
    # Build messages list
    cf_messages = []
    if system_prompt:
        cf_messages.append({"role": "system", "content": system_prompt})
    
    for msg in messages:
        cf_messages.append({
            "role": msg.role,
            "content": msg.content,
        })
    
    # Cloudflare Workers AI endpoint
    model_path = config.model.lstrip('@cf/')
    url = f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/@cf/{model_path}"
    
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    
    async with httpx.AsyncClient(timeout=120.0) as client:
        async with client.stream(
            "POST",
            url,
            headers=headers,
            json={
                "messages": cf_messages,
                "stream": True,
            },
        ) as response:
            if response.status_code != 200:
                body = await response.aread()
                raise ValueError(f"Cloudflare API error: {response.status_code} - {body.decode()}")
            
            async for line in response.aiter_lines():
                if line.startswith("data: "):
                    data_str = line[6:]
                    if data_str == "[DONE]":
                        break
                    try:
                        data = json.loads(data_str)
                        if "response" in data:
                            yield data["response"]
                    except json.JSONDecodeError:
                        continue


async def complete_with_custom(
    messages: list[Message],
    config: ProviderConfig,
    system_prompt: str | None = None,
) -> str:
    """Complete chat using custom OpenAI-compatible endpoint."""
    if not config.endpoint:
        raise ValueError("Custom endpoint URL is required")
    
    # Build messages list (OpenAI format)
    api_messages = []
    if system_prompt:
        api_messages.append({"role": "system", "content": system_prompt})
    
    for msg in messages:
        api_messages.append({
            "role": msg.role,
            "content": msg.content,
        })
    
    headers = {"Content-Type": "application/json"}
    if config.api_key:
        headers["Authorization"] = f"Bearer {config.api_key}"
    
    async with httpx.AsyncClient(timeout=120.0) as client:
        response = await client.post(
            f"{config.endpoint}/v1/chat/completions",
            headers=headers,
            json={
                "model": config.model,
                "messages": api_messages,
                "max_tokens": 4096,
            },
        )
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["message"]["content"]


async def complete_chat(
    messages: list[Message],
    config: ProviderConfig,
    system_prompt: str | None = None,
) -> str:
    """
    Route chat completion to the appropriate provider.
    """
    if not config.enabled:
        raise ValueError(f"Provider '{config.name}' is not enabled")
    
    if config.type == "anthropic":
        return await complete_with_anthropic(messages, config, system_prompt)
    elif config.type == "openai":
        return await complete_with_openai(messages, config, system_prompt)
    elif config.type == "ollama":
        return await complete_with_ollama(messages, config, system_prompt)
    elif config.type == "cloudflare":
        return await complete_with_cloudflare(messages, config, system_prompt)
    elif config.type == "custom":
        return await complete_with_custom(messages, config, system_prompt)
    else:
        raise ValueError(f"Unknown provider type: {config.type}")


# Known multimodal (vision) model patterns
MULTIMODAL_MODEL_PATTERNS = [
    "qwen3-vl", "qwen2.5vl", "qwen2-vl",
    "llava", "llava-phi", "llava-llama",
    "moondream", "bakllava",
    "gemma3",  # Gemma 3 has vision
    "minicpm-v",
    "llama3.2-vision", "llama4",
    "mistral-small3.1", "mistral-small3.2",  # Vision variants
]


def is_multimodal_model(model_name: str) -> bool:
    """Check if a model supports multimodal (vision) input."""
    model_lower = model_name.lower()
    return any(pattern in model_lower for pattern in MULTIMODAL_MODEL_PATTERNS)


# ============================================
# TOOL DEFINITIONS FOR FUNCTION CALLING
# ============================================

AVAILABLE_TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "search_web",
            "description": "Search the web for information using Wikipedia or DuckDuckGo. Use this when you need to find current information, facts, or details about topics you're unsure about.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The search query to look up"
                    },
                    "engine": {
                        "type": "string",
                        "enum": ["wikipedia", "duckduckgo"],
                        "description": "Search engine to use. Wikipedia for encyclopedic content, DuckDuckGo for general web search."
                    }
                },
                "required": ["query"]
            }
        }
    }
]


async def execute_tool(tool_name: str, arguments: dict) -> str:
    """Execute a tool by calling our own local API endpoints."""
    import json
    import logging
    logger = logging.getLogger(__name__)
    
    if tool_name == "search_web":
        query = arguments.get("query", "")
        engine = arguments.get("engine", "wikipedia")
        
        # Call our own search API endpoint (already proven to work)
        try:
            async with httpx.AsyncClient(timeout=15.0) as client:
                response = await client.post(
                    "http://localhost:8000/api/search/web",
                    json={"query": query, "engine": engine, "limit": 5},
                )
                response.raise_for_status()
                data = response.json()
                results = data.get("results", [])
                
                if results:
                    # Format results for the LLM context
                    formatted = []
                    for r in results[:5]:
                        entry = f"- **{r.get('title', 'Untitled')}**: {r.get('snippet', '')[:400]}"
                        if r.get("url"):
                            entry += f"\n  Source: {r['url']}"
                        formatted.append(entry)
                    return "\n\n".join(formatted)
                
                return f"No results found for '{query}' on {engine}."
                
        except Exception as e:
            logger.error(f"[execute_tool] search_web failed: {e}")
            return f"Search failed: {str(e)}"
    
    return json.dumps({"error": f"Unknown tool: {tool_name}"})


async def stream_with_ollama_tools(
    messages: list[Message],
    config: ProviderConfig,
    system_prompt: str | None = None,
    tools: list[dict] | None = None,
) -> AsyncGenerator[str, None]:
    """Stream chat with Ollama, supporting tool calls. Handles tool execution loop."""
    import json
    
    endpoint = (config.endpoint or "http://localhost:11434").rstrip('/')
    
    # Build messages list
    ollama_messages = []
    if system_prompt:
        ollama_messages.append({"role": "system", "content": system_prompt})
    
    for msg in messages:
        message_dict: dict = {
            "role": msg.role,
            "content": msg.content,
        }
        if msg.images:
            message_dict["images"] = msg.images
        ollama_messages.append(message_dict)
    
    # First request - may return tool calls
    request_body = {
        "model": config.model,
        "messages": ollama_messages,
        "stream": True,
    }
    
    # Only add tools if provided (not all models support it)
    if tools:
        request_body["tools"] = tools
    
    async with httpx.AsyncClient(timeout=120.0) as client:
        full_response = ""
        tool_calls = []
        
        async with client.stream(
            "POST",
            f"{endpoint}/api/chat",
            json=request_body,
        ) as response:
            response.raise_for_status()
            async for line in response.aiter_lines():
                if line:
                    try:
                        data = json.loads(line)
                        
                        # Check for tool calls
                        if "message" in data:
                            msg = data["message"]
                            if "tool_calls" in msg and msg["tool_calls"]:
                                tool_calls.extend(msg["tool_calls"])
                            if "content" in msg and msg["content"]:
                                full_response += msg["content"]
                                yield msg["content"]
                                
                    except json.JSONDecodeError:
                        continue
        
        # If there were tool calls, execute them and continue
        if tool_calls:
            yield "\n\n*Searching...*\n\n"
            
            # Execute each tool call
            for tool_call in tool_calls:
                func = tool_call.get("function", {})
                tool_name = func.get("name", "")
                arguments = func.get("arguments", {})
                
                # Execute the tool
                tool_result = await execute_tool(tool_name, arguments)
                
                # Add assistant message with tool call and tool response to messages
                ollama_messages.append({
                    "role": "assistant",
                    "content": full_response,
                    "tool_calls": tool_calls,
                })
                ollama_messages.append({
                    "role": "tool",
                    "content": tool_result,
                })
            
            # Continue the conversation with tool results (no tools this time to avoid loop)
            async with client.stream(
                "POST",
                f"{endpoint}/api/chat",
                json={
                    "model": config.model,
                    "messages": ollama_messages,
                    "stream": True,
                },
            ) as response:
                response.raise_for_status()
                async for line in response.aiter_lines():
                    if line:
                        try:
                            data = json.loads(line)
                            if "message" in data and "content" in data["message"]:
                                yield data["message"]["content"]
                        except json.JSONDecodeError:
                            continue


async def check_ollama_status(endpoint: str = "http://localhost:11434") -> dict:
    """Check if Ollama is running and list available models with multimodal info."""
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            response = await client.get(f"{endpoint}/api/tags")
            response.raise_for_status()
            data = response.json()
            models = [m["name"] for m in data.get("models", [])]
            multimodal_models = [m for m in models if is_multimodal_model(m)]
            return {
                "status": "online",
                "models": models,
                "multimodal_models": multimodal_models,
            }
    except Exception as e:
        return {"status": "offline", "error": str(e), "models": [], "multimodal_models": []}


async def stream_with_abi(
    messages: list[Message],
    config: ProviderConfig,
    system_prompt: str | None = None,
) -> AsyncGenerator[str, None]:
    """Stream chat completion from ABI server (custom API).
    
    ABI uses a custom endpoint format:
    POST {endpoint}/agents/{agent_name}/stream-completion?token={token}
    
    Body format (expected):
    {
        "prompt": "user message here",
        "thread_id": "conversation_id"
    }
    """
    import logging
    logger = logging.getLogger(__name__)
    
    endpoint = config.endpoint.rstrip("/")
    agent_name = config.model  # e.g., "ForvisMazars_Corporate"
    token = config.api_key
    
    # ABI expects a single "prompt" field with the latest user message
    # and a thread_id for conversation continuity
    latest_user_message = None
    for msg in reversed(messages):
        if msg.role == "user":
            latest_user_message = msg.content
            break
    
    if not latest_user_message:
        logger.warning("No user message found in conversation")
        yield "Error: No user message to send"
        return
    
    # Build request URL with token as query param
    url = f"{endpoint}/agents/{agent_name}/stream-completion?token={token}"
    
    logger.info(f"🔌 Streaming from ABI: {url} (agent: {agent_name})")
    
    # Use conversation ID as thread_id (or generate one from messages hash)
    thread_id = str(hash(tuple(m.content for m in messages[-5:])))  # Last 5 messages as thread context
    
    request_body = {
        "prompt": latest_user_message,
        "thread_id": thread_id,
    }
    
    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            async with client.stream(
                "POST",
                url,
                json=request_body,
                headers={"Content-Type": "application/json"},
            ) as response:
                response.raise_for_status()
                
                # ABI returns Server-Sent Events (SSE) format:
                # event: ai_message
                # data: content here
                # event: message
                # data: content here (DUPLICATE!)
                #
                # Strategy: Track the current event type and only yield from 'ai_message' events
                # to avoid duplicates. Don't reset current_event until a new event: line arrives.
                current_event = None
                async for line in response.aiter_lines():
                    if not line or not line.strip():
                        continue
                    
                    # Parse SSE format
                    if line.startswith("event: "):
                        current_event = line[7:].strip()  # Extract event name
                    elif line.startswith("data: "):
                        content = line[6:]  # Remove "data: " prefix
                        
                        # Skip [DONE] marker
                        if content == "[DONE]":
                            break
                        
                        # Only yield content from 'ai_message' events to avoid duplicates
                        # (ABI sends both 'ai_message' and 'message' with same content)
                        # DON'T reset current_event here - it stays until next event: line
                        if current_event == "ai_message" and content.strip():
                            logger.debug(f"📨 ABI chunk ({current_event}): {content[:100]}...")
                            yield content
                            
    except httpx.HTTPStatusError as e:
        logger.error(f"ABI API error: {e.response.status_code} - {e.response.text}")
        yield f"\n\n**Error:** ABI server returned {e.response.status_code}\n\n{e.response.text}"
    except Exception as e:
        logger.error(f"ABI streaming error: {e}")
        yield f"\n\n**Error:** Failed to connect to ABI server: {str(e)}"

