from typing import Dict, List, AsyncIterator, TypedDict, Literal, Optional, Any
import time
import json
import logging

from openai import AsyncOpenAI

from animuz_core.genai.base import BaseAgent
from animuz_core.genai.tools import OpenAITool, MCPTools

logger = logging.getLogger(__name__)


# Type definitions for output structure
class MessageOutput(TypedDict):
    type: Literal["message"]
    role: str
    content: str

class FunctionCallOutput(TypedDict):
    type: Literal["function_call"]
    name: str
    arguments: str
    call_id: str

class FunctionCallResultOutput(TypedDict, total=False):
    type: Literal["function_call_output"]
    call_id: str
    output: str
    metadata: Any  # Filtered out before sending to OpenAI

OutputItem = MessageOutput | FunctionCallOutput | FunctionCallResultOutput


class OpenAIAgentClientStreaming(BaseAgent):
    """Streaming OpenAI Agent Client that yields SSE-compatible events.

    Supports both local tool dispatch and MCP-based tool execution.
    Uses the OpenAI Responses API for streaming.

    Args:
        user_chat_id: User/tenant ID for multi-tenant tool calls
        tools: Local tool dict (name -> (schema, callable))
        model: OpenAI model name (default: gpt-4o-mini)
        max_token: Max output tokens
        mcp_url: MCP server URL. If provided, tools are fetched from MCP.
        user_context: User context dict passed to MCP tool executions
        mcp_api_key: API key for MCP server. Defaults to MCP_API_KEY env var.
    """
    def __init__(
        self,
        user_chat_id: str,
        tools: Dict[str, tuple] = None,
        model: str = None,
        max_token: int = 2000,
        mcp_url: str = None,
        user_context: Dict = None,
        mcp_api_key: str = None,
        profiler=None,
    ) -> None:
        self.client = AsyncOpenAI()
        self.temperature = 0.0
        self.model = model if model else "gpt-4o-mini"
        self.max_token = max_token
        self.profiler = profiler
        self.mcp = None
        if mcp_url:
            self.mcp = MCPTools(mcp_url, user_chat_id, user_context, mcp_api_key, profiler=profiler)
        elif tools:
            self.add_tools(tools)

    def add_tools(self, tools: Dict[str, tuple]):
        self.tools = tools
        self.tool_dump = [
            {"type": "function", "function": tool[0].model_dump()}
            for tool in self.tools.values()
        ]

    def _filter_for_openai(self, messages: List[Dict]) -> List[Dict]:
        """Remove non-OpenAI fields (metadata, complete) before sending to API."""
        clean_messages = []
        for msg in messages:
            clean_msg = msg.copy()
            clean_msg.pop("metadata", None)
            clean_msg.pop("complete", None)

            if clean_msg.get("type") == "message":
                clean_messages.append({
                    "type": "message",
                    "role": clean_msg.get("role"),
                    "content": clean_msg.get("content")
                })
            elif clean_msg.get("type") == "function_call":
                clean_messages.append({
                    "type": "function_call",
                    "name": clean_msg.get("name"),
                    "arguments": clean_msg.get("arguments"),
                    "call_id": clean_msg.get("call_id")
                })
            elif clean_msg.get("type") == "function_call_output":
                clean_messages.append({
                    "type": "function_call_output",
                    "call_id": clean_msg.get("call_id"),
                    "output": clean_msg.get("output")
                })
            else:
                clean_messages.append(clean_msg)

        return clean_messages

    async def stream_reply_frontend(
        self,
        messages: List[Dict[str, str]],
        system_prompt: str = None
    ) -> AsyncIterator[dict]:
        """Stream reply events for frontend consumption.

        Yields dict events in two categories:

        STREAMING EVENTS (for real-time display):
            - {"type": "content_chunk", "content": "..."}
            - {"type": "tool_status", "message": "...", "function_name": "..."}
            - {"type": "tool_executing", "message": "...", "function_name": "...",
               "function_call_id": "...", "function_arguments": {...}}

        COMPLETE EVENTS (for backend/DB storage, marked with "complete": True):
            - {"type": "message", "role": "assistant", "content": "...", "complete": True}
            - {"type": "function_call", "name": "...", "arguments": "...", "call_id": "...", "complete": True}
            - {"type": "function_call_output", "call_id": "...", "output": "...", "complete": True}
        """
        messages_with_prompt = [{"role": "system", "content": system_prompt}] + messages
        user_message = messages[-1]["content"]

        output: List[OutputItem] = []

        async for event in self._create_message_streaming(
            self._filter_for_openai(messages_with_prompt),
            output
        ):
            yield event

        # Tool calling loop
        while self._has_function_calls_in_output(output):
            function_calls = self._extract_function_calls_from_output(output)

            for func_call in function_calls:
                function_name = func_call.get("name")
                args = json.loads(func_call.get("arguments", "{}"))
                call_id = func_call.get("call_id")

                logger.info(f"CALLING FUNCTION: {function_name} WITH ARGS: {str(args)}")

                yield {
                    "type": "tool_executing",
                    "message": f"Executing {function_name}...",
                    "function_name": function_name,
                    "function_call_id": call_id,
                    "function_arguments": args
                }

                result, metadata = await self.process_tool_use_mcp(function_name, args, user_message)

                tool_result: FunctionCallResultOutput = {
                    "type": "function_call_output",
                    "call_id": call_id,
                    "output": result,
                    "metadata": metadata
                }
                output.append(tool_result)
                messages_with_prompt.append(tool_result)

                yield {**tool_result, "complete": True}

            async for event in self._create_message_streaming(
                self._filter_for_openai(messages_with_prompt),
                output
            ):
                yield event

    async def _create_message_streaming(
        self,
        history: List[Dict],
        output: List[OutputItem]
    ) -> AsyncIterator[dict]:
        """Stream a single OpenAI Responses API call and accumulate output."""
        if self.profiler:
            self.profiler.mark("openai_call_start")
        start = time.time()

        # GPT-5 models don't support temperature
        create_kwargs = dict(
            model=self.model,
            max_output_tokens=self.max_token,
            input=history,
            tools=self.mcp.tools if self.mcp else getattr(self, 'tool_dump', []),
            stream=True
        )
        if self.model not in ["gpt-5", "gpt-5-mini", "gpt-5-nano"]:
            create_kwargs["temperature"] = self.temperature

        stream = await self.client.responses.create(**create_kwargs)

        accumulated_content = ""
        final_tool_calls = {}

        async for chunk in stream:
            if chunk.type == "response.in_progress":
                yield {"type": "in_progress"}
            elif chunk.type == "response.completed":
                yield {"type": "completed"}
            elif chunk.type == "response.error":
                raise Exception(f"OpenAI streaming error: {chunk.error.message}")

            if chunk.type == "response.output_text.delta":
                text_chunk = chunk.delta
                accumulated_content += text_chunk
                yield {"type": "content_chunk", "content": text_chunk}

            elif chunk.type == "response.output_item.added":
                if chunk.item.type == "function_call":
                    index = chunk.output_index
                    final_tool_calls[index] = {
                        "function_name": chunk.item.name,
                        "function_call_id": chunk.item.call_id,
                        "arguments": ""
                    }
                    yield {
                        "type": "tool_status",
                        "message": f"Preparing to use {chunk.item.name}...",
                        "function_name": chunk.item.name,
                        "function_call_id": chunk.item.call_id
                    }

            elif chunk.type == "response.function_call_arguments.delta":
                index = chunk.output_index
                if index in final_tool_calls:
                    final_tool_calls[index]["arguments"] += chunk.delta

        openai_duration = time.time() - start
        logger.info(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": "OpenAI API latency (streaming)",
            "latency": round(openai_duration, 4),
        }))
        if self.profiler:
            self.profiler.mark("openai_call_end")
            self.profiler.record_duration("openai_api", openai_duration)

        if accumulated_content:
            message_output: MessageOutput = {
                "type": "message",
                "role": "assistant",
                "content": accumulated_content
            }
            output.append(message_output)
            yield {**message_output, "complete": True}

        for call in final_tool_calls.values():
            function_call_output: FunctionCallOutput = {
                "type": "function_call",
                "name": call["function_name"],
                "arguments": call["arguments"],
                "call_id": call["function_call_id"]
            }
            output.append(function_call_output)
            yield {**function_call_output, "complete": True}

    async def process_tool_use_mcp(self, tool_name: str, tool_input: dict, user_message: str):
        """Execute tool via MCP and return (result, metadata)."""
        start = time.time()
        try:
            response = self.mcp.execute_tool(tool_name, user_message=user_message, **tool_input)
            result, metadata = self.mcp.parse_response(tool_name, response, tool_input)

            logger.info(json.dumps({
                "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
                "message": f"Tool {tool_name} used successfully",
                "latency": round(time.time() - start, 4),
            }, ensure_ascii=False))

            return result, metadata
        except Exception as e:
            logger.error(f"Tool {tool_name} error: {str(e)}")
            return "ERROR", None

    def _has_function_calls_in_output(self, output: List[OutputItem]) -> bool:
        """Check if the last batch of output items contains function calls."""
        if not output:
            return False
        for item in reversed(output):
            if item.get("type") == "function_call":
                return True
            elif item.get("type") == "function_call_output":
                return False
        return False

    def _extract_function_calls_from_output(self, output: List[OutputItem]) -> List[Dict[str, str]]:
        """Extract function calls from the last batch of output items."""
        if not output:
            return []
        last_batch = []
        for item in reversed(output):
            if item.get("type") == "function_call_output":
                break
            last_batch.insert(0, item)
        return [
            {"name": item.get("name", ""), "arguments": item.get("arguments", "{}"), "call_id": item.get("call_id", "")}
            for item in last_batch
            if item.get("type") == "function_call"
        ]

    async def get_reply_frontend(self, messages: List[Dict[str, str]], system_prompt: str = None) -> str:
        raise NotImplementedError("This client is streaming-only. Use stream_reply_frontend instead.")
