"""
Custom Provider Module

Generic OpenAI-compatible provider for any API endpoint:
- Ollama, LM Studio, Together, Groq, Fireworks, etc.
- Local models via compatible APIs
- Custom enterprise deployments
"""

import json
import time
from typing import Optional

import httpx

from ai_coder.llm.interface import (
    LLMProvider,
    LLMConfig,
    Message,
    ToolCall,
    CompletionResponse,
    LLMError,
    RateLimitError,
    AuthenticationError,
    register_provider,
)


@register_provider("custom")
class CustomProvider(LLMProvider):
    """
    Generic OpenAI-compatible API provider.
    
    Works with any API that follows OpenAI's chat completion format:
    - Ollama (http://localhost:11434/v1)
    - LM Studio (http://localhost:1234/v1)
    - Together AI (https://api.together.xyz/v1)
    - Groq (https://api.groq.com/openai/v1)
    - Fireworks (https://api.fireworks.ai/inference/v1)
    - Azure OpenAI
    - Any OpenAI-compatible endpoint
    """

    def __init__(self, config: LLMConfig):
        super().__init__(config)
        
        # Determine base URL
        self.base_url = config.api_base or "http://localhost:11434/v1"
        self.base_url = self.base_url.rstrip("/")
        
        # Setup headers
        self.headers = {
            "Content-Type": "application/json",
        }
        if config.api_key:
            self.headers["Authorization"] = f"Bearer {config.api_key}"
        
        self.client = httpx.Client(timeout=config.timeout)

    def complete(
        self,
        messages: list[Message],
        tools: Optional[list[dict]] = None,
    ) -> CompletionResponse:
        """Generate completion using OpenAI-compatible API."""
        formatted_messages = self._format_messages(messages)

        payload = {
            "model": self.config.model,
            "messages": formatted_messages,
            "max_tokens": self.config.max_tokens,
            "temperature": self.config.temperature,
        }

        # Only add tools if provided and non-empty
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = "auto"

        url = f"{self.base_url}/chat/completions"

        for attempt in range(self.config.max_retries):
            try:
                response = self.client.post(url, headers=self.headers, json=payload)
                
                if response.status_code == 401:
                    raise AuthenticationError("API authentication failed")
                
                if response.status_code == 429:
                    if attempt < self.config.max_retries - 1:
                        wait_time = 2 ** attempt
                        time.sleep(wait_time)
                        continue
                    raise RateLimitError("Rate limit exceeded")
                
                if response.status_code >= 400:
                    error_text = response.text[:500]
                    raise LLMError(f"API error ({response.status_code}): {error_text}")
                
                return self._parse_response(response.json())

            except httpx.TimeoutException:
                if attempt < self.config.max_retries - 1:
                    time.sleep(1)
                    continue
                raise LLMError("Request timed out")
            
            except httpx.RequestError as e:
                if attempt < self.config.max_retries - 1:
                    time.sleep(1)
                    continue
                raise LLMError(f"Connection error: {e}")

        raise LLMError("Max retries exceeded")

    def _format_messages(self, messages: list[Message]) -> list[dict]:
        """Format messages for OpenAI-compatible API."""
        formatted = []

        for msg in messages:
            if msg.role == "tool":
                formatted.append({
                    "role": "tool",
                    "content": msg.content,
                    "tool_call_id": msg.tool_call_id or "",
                })
            elif msg.tool_calls:
                formatted.append({
                    "role": "assistant",
                    "content": msg.content or "",
                    "tool_calls": msg.tool_calls,
                })
            else:
                formatted.append({
                    "role": msg.role,
                    "content": msg.content,
                })

        return formatted

    def _parse_response(self, data: dict) -> CompletionResponse:
        """Parse API response."""
        choices = data.get("choices", [])
        if not choices:
            return CompletionResponse(content="No response generated")

        choice = choices[0]
        message = choice.get("message", {})
        content = message.get("content", "")

        tool_calls = []
        raw_tool_calls = message.get("tool_calls", [])
        
        # Native function calling
        for tc in raw_tool_calls:
            try:
                func = tc.get("function", {})
                args = func.get("arguments", "{}")
                if isinstance(args, str):
                    args = json.loads(args)
                
                tool_calls.append(ToolCall(
                    id=tc.get("id", f"call_{len(tool_calls)}"),
                    name=func.get("name", ""),
                    arguments=args,
                ))
            except json.JSONDecodeError:
                pass

        # Fallback: Parse JSON tool calls from text output
        # This helps models that output JSON instead of using function calling
        if not tool_calls and content:
            tool_call = self._extract_json_tool_call(content)
            if tool_call:
                tool_calls.append(tool_call)
                content = None  # Clear content since we extracted a tool call

        usage = None
        if "usage" in data:
            usage = {
                "prompt_tokens": data["usage"].get("prompt_tokens", 0),
                "completion_tokens": data["usage"].get("completion_tokens", 0),
                "total_tokens": data["usage"].get("total_tokens", 0),
            }

        return CompletionResponse(
            content=content,
            tool_calls=tool_calls,
            finish_reason=choice.get("finish_reason", "stop"),
            usage=usage,
            raw_response=data,
        )

    def _extract_json_tool_call(self, content: str) -> Optional[ToolCall]:
        """Extract a tool call from JSON in text content."""
        import re
        
        # Try to find JSON in code blocks or raw JSON
        patterns = [
            r'```json\s*(\{.*?\})\s*```',  # ```json {...} ```
            r'```\s*(\{.*?\})\s*```',       # ``` {...} ```
            r'(\{[^{}]*"name"[^{}]*"arguments"[^{}]*\})',  # Inline JSON
        ]
        
        for pattern in patterns:
            match = re.search(pattern, content, re.DOTALL)
            if match:
                try:
                    data = json.loads(match.group(1))
                    name = data.get("name", "")
                    args = data.get("arguments", {})
                    
                    if name:
                        if isinstance(args, str):
                            args = json.loads(args)
                        return ToolCall(
                            id=f"extracted_{name}",
                            name=name,
                            arguments=args if isinstance(args, dict) else {},
                        )
                except (json.JSONDecodeError, KeyError):
                    continue
        
        return None


# Also register common provider aliases
@register_provider("ollama")
class OllamaProvider(CustomProvider):
    """Ollama local models."""
    
    def __init__(self, config: LLMConfig):
        if not config.api_base:
            config.api_base = "http://localhost:11434/v1"
        super().__init__(config)


@register_provider("lmstudio")
class LMStudioProvider(CustomProvider):
    """LM Studio local models."""
    
    def __init__(self, config: LLMConfig):
        if not config.api_base:
            config.api_base = "http://localhost:1234/v1"
        super().__init__(config)


@register_provider("together")
class TogetherProvider(CustomProvider):
    """Together AI."""
    
    def __init__(self, config: LLMConfig):
        if not config.api_base:
            config.api_base = "https://api.together.xyz/v1"
        super().__init__(config)


@register_provider("groq")
class GroqProvider(CustomProvider):
    """Groq ultra-fast inference."""
    
    def __init__(self, config: LLMConfig):
        if not config.api_base:
            config.api_base = "https://api.groq.com/openai/v1"
        super().__init__(config)


@register_provider("fireworks")
class FireworksProvider(CustomProvider):
    """Fireworks AI."""
    
    def __init__(self, config: LLMConfig):
        if not config.api_base:
            config.api_base = "https://api.fireworks.ai/inference/v1"
        super().__init__(config)


@register_provider("azure")
class AzureOpenAIProvider(CustomProvider):
    """Azure OpenAI Service."""
    
    def __init__(self, config: LLMConfig):
        super().__init__(config)
        # Azure uses api-key header instead of Bearer token
        if config.api_key:
            self.headers["api-key"] = config.api_key
            self.headers.pop("Authorization", None)
