from typing import Any, AsyncIterator

from agents import Agent, Runner
from agents.extensions.models.litellm_model import LitellmModel

from chatkit.server import ChatKitServer, StreamingResult
from chatkit.store import Store
from chatkit.types import (
    ThreadMetadata, ThreadItem,
    ThreadItemAddedEvent, ThreadItemDoneEvent, ThreadItemUpdatedEvent,
    AssistantMessageItem
)
from chatkit.agents import AgentContext, stream_agent_response, ThreadItemConverter


class ChatKitServerWithMemory(ChatKitServer[dict]):
    """ChatKit server with full conversation memory and LiteLLM ID fix"""

    def __init__(self, data_store: Store, model: LitellmModel, instructions: str):
        super().__init__(data_store)

        self.agent = Agent[AgentContext](
            name="Assistant",
            instructions=instructions,
            model=model,
        )
        # Use ChatKit's official converter for proper item handling
        self.converter = ThreadItemConverter()

    async def respond(
        self,
        thread: ThreadMetadata,
        input: Any,
        context: dict
    ) -> AsyncIterator:
        """Generate response with full conversation context and unique IDs"""

        # Create agent context
        agent_context = AgentContext(
            thread=thread,
            store=self.store,
            request_context=context,
        )

        # Load all thread items from store
        page = await self.store.load_thread_items(
            thread.id,
            after=None,
            limit=100,
            order="asc",
            context=context
        )
        all_items = list(page.data)

        # Add current input to the conversation
        if input:
            all_items.append(input)

        print(f"[Server] Processing {len(all_items)} items for agent")

        # Convert using ChatKit's official converter
        agent_input = await self.converter.to_agent_input(all_items) if all_items else []

        print(f"[Server] Converted to {len(agent_input)} agent input items")

        # Run agent with full history
        result = Runner.run_streamed(
            self.agent,
            agent_input,
            context=agent_context,
        )

        # CRITICAL: Track ID mappings to fix LiteLLM/Gemini ID collision
        id_mapping: dict[str, str] = {}

        async for event in stream_agent_response(agent_context, result):
            # Fix potential ID collisions from LiteLLM/Gemini
            if event.type == "thread.item.added":
                if isinstance(event.item, AssistantMessageItem):
                    old_id = event.item.id
                    # Generate unique ID if we haven't seen this response ID
                    if old_id not in id_mapping:
                        new_id = self.store.generate_item_id("message", thread, context)
                        id_mapping[old_id] = new_id
                        print(f"[Server] Mapping ID {old_id} -> {new_id}")
                    event.item.id = id_mapping[old_id]

            elif event.type == "thread.item.done":
                if isinstance(event.item, AssistantMessageItem):
                    old_id = event.item.id
                    if old_id in id_mapping:
                        event.item.id = id_mapping[old_id]

            elif event.type == "thread.item.updated":
                if event.item_id in id_mapping:
                    event.item_id = id_mapping[event.item_id]

            yield event

from agents.extensions.models.litellm_model import LitellmModel

# Create model (Gemini, OpenAI, Anthropic, etc.)
model = LitellmModel(
    model="gemini/gemini-2.0-flash",
    api_key=os.getenv("GEMINI_API_KEY"),
)

# Create store
store = MemoryStore()

# Create server WITH memory AND ID fix
server = ChatKitServerWithMemory(
    data_store=store,
    model=model,
    instructions="You are a helpful assistant. Remember what the user tells you."
)

yield ThreadItemAddedEvent(
    item=AssistantMessageItem(
        id=item.id,  # <-- Uses provider's ID, may collide!
        ...
    )
)

User: My name is Alice
Assistant: Nice to meet you, Alice!

User: What's my name?
Assistant: Your name is Alice.  <-- CORRECT (has memory, separate messages)

# Without ID fix:
# Both responses might overwrite each other, showing only one message

@app.get("/debug/threads")
async def debug_threads():
    result = {}
    for thread_id, state in store._threads.items():
        items = [{"id": i.id, "type": type(i).__name__} for i in state.items]
        result[thread_id] = {"items": items, "count": len(items)}
    return result