# Copyright 2024 Flyto
# Licensed under the Apache License, Version 2.0
"""OpenAI and OpenAI-compatible provider (function calling loop)."""
import json
import logging
from typing import Any, Dict, List, Optional, Tuple

from flyto_ai.models import StreamCallback, StreamEvent, StreamEventType
from flyto_ai.providers.base import (
    DispatchFn, LLMProvider, dispatch_and_log_tool, fire_stream as _fire,
)

logger = logging.getLogger(__name__)


_BROWSER_INTENT_PATTERNS = [
    "http://", "https://", ".com", ".org", ".net", ".tw", ".io", ".dev",
    "搜尋", "幫我到", "打開", "瀏覽", "上網", "查詢", "查一下", "看看",
    "google", "search", "browse", "website", "scrape", "go to",
    "tixcraft", "wikipedia", "youtube",
]


def _looks_like_browser_task(messages: list) -> bool:
    """Check if the user query requires browser automation."""
    for msg in reversed(messages):
        if msg.get("role") == "user":
            text = msg.get("content", "")
            if isinstance(text, str):
                lower = text.lower()
                return any(p in lower for p in _BROWSER_INTENT_PATTERNS)
            break
    return False


class OpenAIProvider(LLMProvider):
    """OpenAI provider with function calling loop."""

    def __init__(
        self,
        api_key: str,
        model: str = "gpt-4o-mini",
        temperature: float = 0.7,
        max_tokens: int = 4096,
        base_url: Optional[str] = None,
    ) -> None:
        self._api_key = api_key
        self._model = model
        self._temperature = temperature
        self._max_tokens = max_tokens
        self._base_url = base_url
        self._client = None

    def __repr__(self) -> str:
        key_hint = "{}...".format(self._api_key[:4]) if self._api_key and len(self._api_key) > 4 else "***"
        return "OpenAIProvider(model={!r}, api_key={!r})".format(self._model, key_hint)

    def _make_client(self):
        if self._client is None:
            import openai
            kwargs = {"api_key": self._api_key}
            if self._base_url:
                kwargs["base_url"] = self._base_url
            self._client = openai.AsyncOpenAI(**kwargs)
        return self._client

    def _is_native_openai(self) -> bool:
        """True if talking to real OpenAI API (not Ollama / custom base_url)."""
        if not self._base_url:
            return True
        from urllib.parse import urlparse
        host = urlparse(self._base_url).hostname or ""
        return host == "api.openai.com" or host.endswith(".openai.com")

    async def chat(
        self,
        messages: List[Dict[str, Any]],
        system_prompt: str,
        tools: List[Dict],
        dispatch_fn: DispatchFn,
        max_rounds: int = 30,
        on_stream: Optional[StreamCallback] = None,
    ) -> Tuple[str, List[Dict[str, Any]], int, Dict[str, int]]:
        client = self._make_client()

        # Convert to OpenAI function format
        openai_tools = [
            {
                "type": "function",
                "function": {
                    "name": t["name"],
                    "description": t["description"],
                    "parameters": t["inputSchema"],
                },
            }
            for t in tools
        ]

        full_messages = [{"role": "system", "content": system_prompt}]
        full_messages.extend(messages)

        tool_call_log: List[Dict[str, Any]] = []
        total_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}

        for round_num in range(max_rounds):
            create_kwargs = {
                "model": self._model,
                "messages": full_messages,
                "max_tokens": self._max_tokens,
                "temperature": self._temperature,
            }
            if openai_tools:
                create_kwargs["tools"] = openai_tools
                # Force tool use on first round to prevent LLM from
                # skipping tools entirely for browser/web tasks.
                if round_num == 0 and _looks_like_browser_task(messages):
                    create_kwargs["tool_choice"] = "required"
                else:
                    create_kwargs["tool_choice"] = "auto"

            # ── Streaming path ──────────────────────────────────
            if on_stream is not None:
                create_kwargs["stream"] = True
                # Request usage in stream (native OpenAI only)
                if self._is_native_openai():
                    create_kwargs["stream_options"] = {"include_usage": True}

                content_parts: List[str] = []
                collected_tool_calls: Dict[int, Dict[str, Any]] = {}
                finish_reason = None

                stream = await client.chat.completions.create(**create_kwargs)
                async for chunk in stream:
                    # Accumulate usage from stream (last chunk)
                    if hasattr(chunk, "usage") and chunk.usage:
                        total_usage["prompt_tokens"] += chunk.usage.prompt_tokens or 0
                        total_usage["completion_tokens"] += chunk.usage.completion_tokens or 0
                        total_usage["total_tokens"] += chunk.usage.total_tokens or 0

                    delta = chunk.choices[0].delta if chunk.choices else None
                    if delta is None:
                        continue

                    # Finish reason
                    if chunk.choices[0].finish_reason:
                        finish_reason = chunk.choices[0].finish_reason

                    # Text content
                    if delta.content:
                        content_parts.append(delta.content)
                        _fire(on_stream, StreamEvent(
                            type=StreamEventType.TOKEN,
                            content=delta.content,
                        ))

                    # Tool call deltas — accumulate by index
                    if delta.tool_calls:
                        for tc_delta in delta.tool_calls:
                            idx = tc_delta.index
                            if idx not in collected_tool_calls:
                                collected_tool_calls[idx] = {
                                    "id": tc_delta.id or "",
                                    "name": "",
                                    "arguments": "",
                                }
                            entry = collected_tool_calls[idx]
                            if tc_delta.id:
                                entry["id"] = tc_delta.id
                            if tc_delta.function:
                                if tc_delta.function.name:
                                    entry["name"] += tc_delta.function.name
                                if tc_delta.function.arguments:
                                    entry["arguments"] += tc_delta.function.arguments

                # No tool calls → return content
                if not collected_tool_calls:
                    content = "".join(content_parts)
                    if finish_reason == "length":
                        content += "\n\n[Note: Response was truncated due to token limit.]"
                    _fire(on_stream, StreamEvent(type=StreamEventType.DONE))
                    return content, tool_call_log, round_num + 1, total_usage

                # Reconstruct assistant message for conversation history
                tc_list = []
                for idx in sorted(collected_tool_calls.keys()):
                    tc = collected_tool_calls[idx]
                    tc_list.append({
                        "id": tc["id"],
                        "type": "function",
                        "function": {
                            "name": tc["name"],
                            "arguments": tc["arguments"],
                        },
                    })

                # Build a message object compatible with the OpenAI API
                import openai.types.chat as _cht
                assistant_tc_objs = [
                    _cht.ChatCompletionMessageToolCall(
                        id=tc["id"],
                        type="function",
                        function=_cht.chat_completion_message_tool_call.Function(
                            name=tc["function"]["name"],
                            arguments=tc["function"]["arguments"],
                        ),
                    )
                    for tc in tc_list
                ]
                # Strip text content from intermediate rounds with tool calls.
                # Fabricated text alongside tool calls pollutes conversation history.
                assistant_msg = _cht.ChatCompletionMessage(
                    role="assistant",
                    content=None,
                    tool_calls=assistant_tc_objs,
                )
                full_messages.append(assistant_msg)

                # Dispatch each tool call via shared helper
                _ask_user_break = False
                for tc in tc_list:
                    func_name = tc["function"]["name"]
                    try:
                        func_args = json.loads(tc["function"]["arguments"])
                    except json.JSONDecodeError:
                        func_args = {}

                    result_str, log_entry, images = await dispatch_and_log_tool(
                        func_name, func_args, dispatch_fn, round_num, on_stream,
                    )
                    tool_call_log.append(log_entry)

                    full_messages.append({
                        "role": "tool",
                        "tool_call_id": tc["id"],
                        "content": result_str,
                    })

                    # ask_user marker — break tool loop, user input needed
                    if log_entry.get("result", {}).get("__ASK_USER__"):
                        _ask_user_break = True
                        # Must still respond to remaining tool_calls to avoid API error
                        for remaining_tc in tc_list[tc_list.index(tc) + 1:]:
                            full_messages.append({
                                "role": "tool",
                                "tool_call_id": remaining_tc["id"],
                                "content": '{"ok": false, "error": "Execution paused — waiting for user input"}',
                            })
                        break

                    # Inject vision user message for native OpenAI (tool messages don't support images)
                    if images and self._is_native_openai():
                        vision_content = [{"type": "text", "text": "[Screenshot from {}]".format(func_name)}]
                        for img in images:
                            vision_content.append({
                                "type": "image_url",
                                "image_url": {
                                    "url": "data:{};base64,{}".format(img.get("media_type", "image/png"), img["base64"]),
                                    "detail": "low",
                                },
                            })
                        full_messages.append({"role": "user", "content": vision_content})

                if _ask_user_break:
                    text = "I need some information from you before I can continue."
                    _fire(on_stream, StreamEvent(type=StreamEventType.DONE))
                    return text, tool_call_log, round_num + 1, total_usage

                continue  # next round

            # ── Non-streaming path ──────────────────────────────
            response = await client.chat.completions.create(**create_kwargs)

            # Accumulate usage
            if response.usage:
                total_usage["prompt_tokens"] += response.usage.prompt_tokens or 0
                total_usage["completion_tokens"] += response.usage.completion_tokens or 0
                total_usage["total_tokens"] += response.usage.total_tokens or 0

            choice = response.choices[0]

            if not choice.message.tool_calls:
                content = choice.message.content or ""
                if choice.finish_reason == "length":
                    content += "\n\n[Note: Response was truncated due to token limit.]"
                return content, tool_call_log, round_num + 1, total_usage

            # Strip text content from intermediate rounds to prevent fabrication.
            # Reconstruct message with content=None so the LLM generates text
            # only AFTER seeing actual tool results.
            import openai.types.chat as _cht2
            stripped_tc_objs = [
                _cht2.ChatCompletionMessageToolCall(
                    id=tc.id,
                    type="function",
                    function=_cht2.chat_completion_message_tool_call.Function(
                        name=tc.function.name,
                        arguments=tc.function.arguments,
                    ),
                )
                for tc in choice.message.tool_calls
            ]
            stripped_msg = _cht2.ChatCompletionMessage(
                role="assistant",
                content=None,
                tool_calls=stripped_tc_objs,
            )
            full_messages.append(stripped_msg)

            _ask_user_break = False
            for tc in choice.message.tool_calls:
                func_name = tc.function.name
                try:
                    func_args = json.loads(tc.function.arguments)
                except json.JSONDecodeError:
                    func_args = {}

                result_str, log_entry, images = await dispatch_and_log_tool(
                    func_name, func_args, dispatch_fn, round_num,
                )
                tool_call_log.append(log_entry)

                full_messages.append({
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": result_str,
                })

                # ask_user marker — break tool loop, user input needed
                if log_entry.get("result", {}).get("__ASK_USER__"):
                    _ask_user_break = True
                    tc_list = list(choice.message.tool_calls)
                    tc_idx = tc_list.index(tc)
                    for remaining in tc_list[tc_idx + 1:]:
                        full_messages.append({
                            "role": "tool",
                            "tool_call_id": remaining.id,
                            "content": '{"ok": false, "error": "Execution paused — waiting for user input"}',
                        })
                    break

                # Inject vision user message for native OpenAI (tool messages don't support images)
                if images and self._is_native_openai():
                    vision_content = [{"type": "text", "text": "[Screenshot from {}]".format(func_name)}]
                    for img in images:
                        vision_content.append({
                            "type": "image_url",
                            "image_url": {
                                "url": "data:{};base64,{}".format(img.get("media_type", "image/png"), img["base64"]),
                                "detail": "low",
                            },
                        })
                    full_messages.append({"role": "user", "content": vision_content})

            if _ask_user_break:
                text = "I need some information from you before I can continue."
                return text, tool_call_log, round_num + 1, total_usage

        # Hit max rounds — force summary with improved prompt
        completed = [tc["function"] for tc in tool_call_log if tc.get("ok", True)]
        failed = [tc["function"] for tc in tool_call_log if not tc.get("ok", True)]
        summary_parts = [
            "You have used all {} tool rounds.".format(max_rounds),
            "Completed: {}".format(", ".join(completed[-5:]) if completed else "none"),
        ]
        if failed:
            summary_parts.append("Failed: {}".format(", ".join(failed[-3:])))
        summary_parts.append(
            "Please summarize what was accomplished, what remains incomplete, "
            "and suggest next steps."
        )
        full_messages.append({
            "role": "user",
            "content": " ".join(summary_parts),
        })
        response = await client.chat.completions.create(
            model=self._model,
            messages=full_messages,
            max_tokens=self._max_tokens,
            temperature=self._temperature,
        )
        if response.usage:
            total_usage["prompt_tokens"] += response.usage.prompt_tokens or 0
            total_usage["completion_tokens"] += response.usage.completion_tokens or 0
            total_usage["total_tokens"] += response.usage.total_tokens or 0

        _fire(on_stream, StreamEvent(type=StreamEventType.DONE))
        return response.choices[0].message.content or "", tool_call_log, max_rounds, total_usage
