from __future__ import annotations
from typing import AsyncIterable, List, Dict
import requests
from fastapi_poe import PoeBot, QueryRequest, PartialResponse, SettingsRequest, SettingsResponse, make_app
from fastapi import Header, HTTPException
from modal import Image, App, asgi_app
import os

# LM Studio endpoint configurations (reuse ngrok URL)
NGROK_URL = "https://fca7-2601-2c1-280-1320-b881-aac7-9186-9365.ngrok-free.app"
LM_STUDIO_CHAT_URL = f"{NGROK_URL}/api/v0/chat/completions"
LM_STUDIO_EMBEDDINGS_URL = f"{NGROK_URL}/api/v0/embeddings"

# Hardcoded model name for LM Studio
MODEL_NAME = "bartowski/Qwen2.5-Coder-32B-Instruct-GGUF/Qwen2.5-Coder-32B-Instruct-IQ2_M.gguf"

# Poe bot access key for the new bot
NEW_BOT_ACCESS_KEY = "Kskn3sHdJh7nkVWGcpwcGOIHYSrsML3l"

# Dictionary to track user-specific conversation history
user_histories: Dict[str, List[dict]] = {}

class main(PoeBot):
    async def get_response(
        self, request: QueryRequest, authorization: str = Header(...)
    ) -> AsyncIterable[PartialResponse]:
        """
        Handle user queries dynamically while validating the Poe access key.
        """
        # Validate the Poe access key
        if authorization != NEW_BOT_ACCESS_KEY:
            print("Warning: Unauthorized access key used.")
            # Proceed without raising an error for testing purposes

        # Extract user identifier
        user_id = request.user_id if hasattr(request, "user_id") else str(hash(request))
        if not user_id:
            yield PartialResponse(text="Error: User identifier not provided.")
            return

        # Extract user message
        user_message = request.query[-1].content  # Latest user input

        # Get or create user-specific conversation history
        if user_id not in user_histories:
            user_histories[user_id] = []
        conversation_history = user_histories[user_id]

        try:
            # 1. Update conversation history with the user message
            conversation_history.append({"role": "user", "content": user_message})

            # 2. Generate response based on conversation history
            response_text = self.get_chat_completion(conversation_history)

            # 3. Add bot response to conversation history
            conversation_history.append({"role": "assistant", "content": response_text})

            # 4. Yield the response
            yield PartialResponse(text=response_text.strip())
        except Exception as e:
            # Graceful error handling
            yield PartialResponse(text=f"An error occurred: {e}")

    def get_chat_completion(self, messages: List[dict]) -> str:
        """
        Send a chat completion request to LM Studio's /api/v0/chat/completions endpoint.
        """
        payload = {
            "model": MODEL_NAME,
            "messages": messages,
            "temperature": 0.7,  # Adjusted temperature for different bot behavior
            "max_tokens": 1024,  # Increased max tokens
            "stream": False  # Do not use streaming for now
        }
        response = requests.post(LM_STUDIO_CHAT_URL, json=payload, timeout=30)  # Increased timeout
        response.raise_for_status()
        generated_text = response.json().get("choices", [{}])[0].get("message", {}).get("content", "")

        # Fallback in case of empty content
        if not generated_text:
            generated_text = "I'm sorry, I couldn't generate a response. Could you please try again?"
        
        return generated_text

    def get_embeddings(self, text: str) -> str:
        """
        Generate text embeddings using LM Studio's /api/v0/embeddings endpoint.
        """
        payload = {"model": "openai/clip-vit-base-patch32", "input": text}  # Update model if needed
        response = requests.post(LM_STUDIO_EMBEDDINGS_URL, json=payload, timeout=10)
        response.raise_for_status()
        embedding = response.json().get("data", [{}])[0].get("embedding", [])
        return f"Embeddings: {embedding}"

    async def get_settings(self, setting: SettingsRequest) -> SettingsResponse:
        """
        Configure the bot's capabilities for Poe, such as enabling attachments.
        """
        return SettingsResponse(
            allow_attachments=False,  # Disable file uploads for now
        )

# Modal configuration for the new bot
REQUIREMENTS = ["fastapi-poe==0.0.24", "requests==2.31.0"]
image = Image.debian_slim().pip_install(REQUIREMENTS)
new_app = App("another-secure-lmstudio-poe-bot")

@new_app.function(image=image)
@asgi_app()
def fastapi_app():
    bot = AnotherSecureLMStudioBot()
    return make_app(bot, allow_without_key=True)