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

from __future__ import annotations

import asyncio
import os
import re
import uuid
from pathlib import Path
from typing import Any

import json_repair
from openai import AsyncOpenAI

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


class _CrossProcessFileLock:
    """Simple POSIX file lock for serializing local model requests across processes."""

    def __init__(self, path: str) -> None:
        self.path = Path(path)
        self._handle = None

    def acquire(self) -> None:
        import fcntl

        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._handle = self.path.open("a+")
        fcntl.flock(self._handle.fileno(), fcntl.LOCK_EX)

    def release(self) -> None:
        import fcntl

        if self._handle is None:
            return
        try:
            fcntl.flock(self._handle.fileno(), fcntl.LOCK_UN)
        finally:
            self._handle.close()
            self._handle = None


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
        self._lock_path = None
        if str(os.environ.get("MADPOD_AGENT_SERIALIZE_LLM_REQUESTS", "")).strip().lower() in {"1", "true", "yes", "on"}:
            self._lock_path = os.environ.get("MADPOD_AGENT_LLM_LOCK_PATH", "/tmp/madpod-agent-llm.lock")
        timeout_s = float(os.environ.get("MADPOD_AGENT_LLM_TIMEOUT_S", "120"))
        # Keep affinity stable for this provider instance to improve backend cache locality.
        self._client = AsyncOpenAI(
            api_key=api_key,
            base_url=api_base,
            default_headers={"x-session-affinity": uuid.uuid4().hex},
            timeout=timeout_s,
            max_retries=0,
        )

    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,
                   reasoning_effort: str | None = None) -> LLMResponse:
        kwargs: dict[str, Any] = {
            "model": model or self.default_model,
            "messages": self._sanitize_empty_content(messages),
            "max_tokens": max(1, max_tokens),
            "temperature": temperature,
        }
        if reasoning_effort:
            kwargs["reasoning_effort"] = reasoning_effort
        if tools:
            kwargs.update(tools=tools, tool_choice="auto")
        lock = _CrossProcessFileLock(self._lock_path) if self._lock_path else None
        try:
            if lock is not None:
                await asyncio.to_thread(lock.acquire)
            return self._parse(await self._client.chat.completions.create(**kwargs))
        except Exception as e:
            return LLMResponse(content=f"Error: {e}", finish_reason="error")
        finally:
            if lock is not None:
                await asyncio.to_thread(lock.release)

    def _parse(self, response: Any) -> LLMResponse:
        choice = response.choices[0]
        msg = choice.message
        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 [])
        ]
        content = msg.content
        if not tool_calls:
            content, tool_calls = _extract_tool_calls_from_content(content)
        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) or None,
        )

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


def _extract_tool_calls_from_content(
    content: str | None,
) -> tuple[str | None, list[ToolCallRequest]]:
    stripped = str(content or "").strip()
    if not stripped:
        return content, []

    for candidate in _content_tool_call_candidates(stripped):
        try:
            parsed = json_repair.loads(candidate)
        except Exception:
            continue
        tool_calls = _tool_calls_from_parsed_payload(parsed)
        if not tool_calls:
            continue

        remaining = stripped.replace(candidate, "", 1).strip()
        if (
            not remaining
            or "Current best provisional answer:" in remaining
            or remaining.startswith("I investigated this")
        ):
            remaining = ""
        return remaining or None, tool_calls

    return content, []


def _content_tool_call_candidates(content: str) -> list[str]:
    candidates: list[str] = []

    def _append(candidate: str) -> None:
        cleaned = str(candidate or "").strip()
        if cleaned and cleaned not in candidates:
            candidates.append(cleaned)

    _append(content)

    for block in re.findall(r"```(?:json)?\s*([\s\S]*?)```", content, flags=re.IGNORECASE):
        _append(block)

    for match in reversed(list(re.finditer(r"\{", content))):
        _append(content[match.start() :])

    return candidates


def _tool_calls_from_parsed_payload(payload: Any) -> list[ToolCallRequest]:
    if isinstance(payload, dict):
        maybe_call = _tool_call_from_mapping(payload)
        return [maybe_call] if maybe_call is not None else []

    if isinstance(payload, list):
        calls = [_tool_call_from_mapping(item) for item in payload]
        return [call for call in calls if call is not None]

    return []


def _tool_call_from_mapping(payload: Any) -> ToolCallRequest | None:
    if not isinstance(payload, dict):
        return None

    name = str(
        payload.get("name")
        or payload.get("tool")
        or payload.get("tool_name")
        or ""
    ).strip()
    if not name:
        return None

    arguments = payload.get("arguments")
    if arguments is None:
        arguments = payload.get("args")
    if arguments is None:
        arguments = payload.get("parameters")

    if isinstance(arguments, str):
        try:
            arguments = json_repair.loads(arguments)
        except Exception:
            return None
    if arguments is None:
        arguments = {}
    if not isinstance(arguments, dict):
        return None

    return ToolCallRequest(
        id=str(payload.get("id") or f"fallback_{uuid.uuid4().hex[:8]}"),
        name=name,
        arguments=arguments,
    )
