"""Direct OpenAI-compatible provider — bypasses LiteLLM."""

from __future__ import annotations

from typing import Any

import json_repair
from openai import AsyncOpenAI

from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest


class CustomProvider(LLMProvider):

    def __init__(self, api_key: str = "no-key", api_base: str = "http://localhost:8000/v1", default_model: str = "default"):
        super().__init__(api_key, api_base)
        self.default_model = default_model
        # Set longer timeout for Ollama (especially with GPU, first load can be slow)
        from openai import Timeout
        # Set separate timeouts: connect, read, write, pool
        # Ollama can be slow, especially on first request
        self._client = AsyncOpenAI(
            api_key=api_key, 
            base_url=api_base,
            timeout=Timeout(
                connect=60.0,  # Connection timeout
                read=600.0,    # Read timeout (10 min for slow Ollama responses)
                write=60.0,    # Write timeout
                pool=60.0      # Pool timeout
            )
        )

    async def chat(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None,
                   model: str | None = None, max_tokens: int = 4096, temperature: float = 0.7) -> LLMResponse:
        kwargs: dict[str, Any] = {"model": model or self.default_model, "messages": messages,
                                  "max_tokens": max(1, max_tokens), "temperature": temperature}
        if tools:
            kwargs.update(tools=tools, tool_choice="auto")
        try:
            import asyncio
            # Add explicit timeout wrapper (longer for Ollama)
            return self._parse(await asyncio.wait_for(
                self._client.chat.completions.create(**kwargs),
                timeout=310.0  # Slightly longer than client timeout (300s)
            ))
        except asyncio.TimeoutError:
            return LLMResponse(content="Error: Request timed out after 310 seconds", finish_reason="error")
        except Exception as e:
            return LLMResponse(content=f"Error: {e}", finish_reason="error")

    def _parse(self, response: Any) -> LLMResponse:
        choice = response.choices[0]
        msg = choice.message
        
        # First, try to get structured tool calls
        tool_calls = [
            ToolCallRequest(id=tc.id, name=tc.function.name,
                            arguments=json_repair.loads(tc.function.arguments) if isinstance(tc.function.arguments, str) else tc.function.arguments)
            for tc in (msg.tool_calls or [])
        ]
        
        # If no structured tool calls, try to parse from content (some OpenAI-compatible backends return JSON in content)
        # Only parse if content looks like it contains a tool call JSON (to avoid false positives)
        content = msg.content or ""
        stripped = content.strip()
        # Note: This list should match tools registered in AgentLoop._register_default_tools().
        # MCP tools are registered dynamically and are prefixed with "mcp_" (allow those too).
        valid_tools = [
            # File tools
            "read_file", "write_file", "edit_file", "list_dir",
            # Shell tool
            "exec",
            # Web tools
            "web_search", "web_fetch",
            # Communication tools
            "message", "spawn",
            # Calendar tool
            "calendar",
            # Cron tool
            "cron",
            # Email tool
            "email",
        ]
        # Check for standard format: {"name": "...", "parameters": {...}}
        has_standard_format = '"name"' in content and '"parameters"' in content
        # Check for calendar tool format: {"action": "...", ...}
        has_calendar_format = '"action"' in content and ("calendar" in content.lower() or any(action in content for action in ["list_events", "create_event", "update_event", "delete_event"]))
        
        # Some backends will return *only* a JSON object as the entire message content.
        # If it looks like a JSON object, attempt parsing even if our heuristics missed it.
        looks_like_json_object = stripped.startswith("{") and stripped.endswith("}")

        if not tool_calls and content and (has_standard_format or has_calendar_format or looks_like_json_object):
            import re
            # Look for JSON tool call patterns: {"name": "exec", "parameters": {...}} or {"action": "list_events", ...}
            # Find complete JSON objects by matching braces
            # Try "action" pattern first (for calendar tool), then "name" pattern
            patterns = [
                (r'\{\s*"action"\s*:\s*"(\w+)"', "action"),  # Calendar tool format
                (r'\{\s*"name"\s*:\s*"(\w+)"', "name"),     # Standard format
            ]
            start_pos = 0
            max_iterations = 10  # Increased for multiple patterns
            iteration = 0
            while iteration < max_iterations:
                iteration += 1
                match = None
                pattern_type = None
                for pattern, ptype in patterns:
                    match = re.search(pattern, content[start_pos:])
                    if match:
                        pattern_type = ptype
                        break
                if not match:
                    break
                
                json_start = start_pos + match.start()
                key_value = match.group(1)
                
                # Find the matching closing brace by counting braces
                brace_count = 0
                json_end = json_start
                found_end = False
                for i, char in enumerate(content[json_start:], json_start):
                    if char == '{':
                        brace_count += 1
                    elif char == '}':
                        brace_count -= 1
                        if brace_count == 0:
                            json_end = i + 1
                            found_end = True
                            break
                
                if found_end:
                    # Try to parse the complete JSON object
                    try:
                        json_str = content[json_start:json_end]
                        tool_obj = json_repair.loads(json_str)
                        
                        # Handle calendar tool format: {"action": "...", ...}
                        if isinstance(tool_obj, dict) and "action" in tool_obj:
                            # This is a calendar tool call in JSON format
                            action = tool_obj.get("action")
                            if action and action in ["list_events", "create_event", "update_event", "delete_event", "delete_events", "check_availability"]:
                                # Convert to calendar tool call format
                                tool_calls.append(ToolCallRequest(
                                    id=f"call_{len(tool_calls)}",
                                    name="calendar",
                                    arguments=tool_obj  # Pass the whole object as arguments
                                ))
                                # Remove the tool call from content
                                content = content[:json_start] + content[json_end:].strip()
                                start_pos = json_start  # Stay at same position since we removed text
                                continue
                        
                        # Handle standard format: {"name": "...", "parameters": {...}}
                        if (isinstance(tool_obj, dict) and 
                            "name" in tool_obj and 
                            "parameters" in tool_obj and
                            isinstance(tool_obj["name"], str) and
                            (tool_obj["name"] in valid_tools or tool_obj["name"].startswith("mcp_"))):
                            tool_calls.append(ToolCallRequest(
                                id=f"call_{len(tool_calls)}",
                                name=tool_obj["name"],
                                arguments=tool_obj["parameters"] if isinstance(tool_obj["parameters"], dict) else {"raw": str(tool_obj["parameters"])}
                            ))
                            # Remove the tool call from content
                            content = content[:json_start] + content[json_end:].strip()
                            start_pos = json_start  # Stay at same position since we removed text
                            continue
                    except Exception:
                        pass  # If parsing fails, skip this match
                
                start_pos = json_start + 1  # Move past this match

            # If we still didn't match embedded objects, try parsing the whole message as a single tool-call JSON object.
            if not tool_calls and looks_like_json_object:
                try:
                    tool_obj = json_repair.loads(stripped)
                    if isinstance(tool_obj, dict) and "action" in tool_obj:
                        action = tool_obj.get("action")
                        if action and action in ["list_events", "create_event", "update_event", "delete_event", "delete_events", "check_availability"]:
                            tool_calls.append(ToolCallRequest(
                                id="call_0",
                                name="calendar",
                                arguments=tool_obj,
                            ))
                            content = ""
                    if isinstance(tool_obj, dict) and "name" in tool_obj and "parameters" in tool_obj:
                        if isinstance(tool_obj["name"], str) and (
                            tool_obj["name"] in valid_tools or tool_obj["name"].startswith("mcp_")
                        ):
                            tool_calls.append(ToolCallRequest(
                                id="call_0",
                                name=tool_obj["name"],
                                arguments=tool_obj["parameters"] if isinstance(tool_obj["parameters"], dict) else {"raw": str(tool_obj["parameters"])},
                            ))
                            content = ""
                except Exception:
                    pass
        
        u = response.usage
        return LLMResponse(
            content=content, tool_calls=tool_calls, finish_reason=choice.finish_reason or "stop",
            usage={"prompt_tokens": u.prompt_tokens, "completion_tokens": u.completion_tokens, "total_tokens": u.total_tokens} if u else {},
            reasoning_content=getattr(msg, "reasoning_content", None),
        )

    def get_default_model(self) -> str:
        return self.default_model
