import os
import time
import asyncio
import logging
import sqlite3
import tiktoken
from uuid import uuid4
from functools import lru_cache
from typing import Optional, List, Dict, Literal
from fastapi import FastAPI, HTTPException, Depends, Security, BackgroundTasks
from fastapi.security import APIKeyHeader
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from openai import OpenAI

# ============================================================================
# Configuration and Setup
# ============================================================================

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("app.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# FastAPI app setup
app = FastAPI()

# API key configuration
API_KEY_NAME = "X-API-Key"
API_KEY = os.environ.get("CHAT_AUTH_KEY", "default_secret_key")
api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)

# Model definitions
ModelID = Literal[
    "openai/gpt-4o-mini",
    "meta-llama/llama-3-70b-instruct",
    "anthropic/claude-3.5-sonnet",
    "deepseek/deepseek-coder",
    "anthropic/claude-3-haiku",
    "openai/gpt-3.5-turbo-instruct",
    "qwen/qwen-72b-chat",
    "google/gemma-2-27b-it"
]

# Pydantic models
class LLMAgentQueryModel(BaseModel):
    prompt: str = Field(..., description="User's query or prompt")
    system_message: Optional[str] = Field(None, description="Custom system message for the conversation")
    model_id: ModelID = Field(
        default="openai/gpt-4o-mini",
        description="ID of the model to use for response generation"
    )
    conversation_id: Optional[str] = Field(None, description="Unique identifier for the conversation")
    user_id: str = Field(..., description="Unique identifier for the user")

    class Config:
        schema_extra = {
            "example": {
                "prompt": "How do I implement a binary search in Python?",
                "system_message": "You are a helpful coding assistant.",
                "model_id": "meta-llama/llama-3-70b-instruct",
                "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
                "user_id": "user123"
            }
        }

# API key and client setup
@lru_cache()
def get_api_keys():
    logger.info("Loading API keys")
    return {
        "OPENROUTER_API_KEY": f"sk-or-v1-{os.environ['OPENROUTER_API_KEY']}",
    }

api_keys = get_api_keys()
or_client = OpenAI(api_key=api_keys["OPENROUTER_API_KEY"], base_url="https://openrouter.ai/api/v1")

# In-memory storage for conversations
conversations: Dict[str, List[Dict[str, str]]] = {}
last_activity: Dict[str, float] = {}

# Token encoding
encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")

# ============================================================================
# Database Functions
# ============================================================================

DB_PATH = '/app/data/conversations.db'

def init_db():
    logger.info("Initializing database")
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS conversations
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                  user_id TEXT,
                  conversation_id TEXT,
                  message TEXT,
                  response TEXT,
                  timestamp DATETIME DEFAULT CURRENT_TIMESTAMP)''')
    conn.commit()
    conn.close()
    logger.info("Database initialized successfully")

def update_db(user_id, conversation_id, message, response):
    logger.info(f"Updating database for conversation: {conversation_id}")
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute('''INSERT INTO conversations (user_id, conversation_id, message, response)
                 VALUES (?, ?, ?, ?)''', (user_id, conversation_id, message, response))
    conn.commit()
    conn.close()
    logger.info("Database updated successfully")

# ============================================================================
# Utility Functions
# ============================================================================

def calculate_tokens(msgs):
    return sum(len(encoding.encode(str(m))) for m in msgs)

def limit_conversation_history(conversation: List[Dict[str, str]], max_tokens: int = 4000) -> List[Dict[str, str]]:
    """Limit the conversation history to a maximum number of tokens."""
    limited_conversation = []
    current_tokens = 0

    for message in reversed(conversation):
        message_tokens = calculate_tokens([message])
        if current_tokens + message_tokens > max_tokens:
            break
        limited_conversation.insert(0, message)
        current_tokens += message_tokens

    return limited_conversation

async def verify_api_key(api_key: str = Security(api_key_header)):
    if api_key != API_KEY:
        logger.warning("Invalid API key used")
        raise HTTPException(status_code=403, detail="Could not validate credentials")
    return api_key

# ============================================================================
# LLM Interaction Functions
# ============================================================================

def chat_with_llama_stream(messages, model="meta-llama/llama-3-70b-instruct", max_output_tokens=2500):
    logger.info(f"Starting chat with model: {model}")
    try:
        response = or_client.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=max_output_tokens,
            stream=True
        )
        
        full_response = ""
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                content = chunk.choices[0].delta.content
                full_response += content
                yield content
        
        # After streaming, add the full response to the conversation history
        messages.append({"role": "assistant", "content": full_response})
        logger.info("Chat completed successfully")
    except Exception as e:
        logger.error(f"Error in model response: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error in model response: {str(e)}")

# ============================================================================
# Background Tasks
# ============================================================================

async def clear_inactive_conversations():
    while True:
        logger.info("Clearing inactive conversations")
        current_time = time.time()
        inactive_convos = [conv_id for conv_id, last_time in last_activity.items() 
                           if current_time - last_time > 1800]  # 30 minutes
        for conv_id in inactive_convos:
            if conv_id in conversations:
                del conversations[conv_id]
            if conv_id in last_activity:
                del last_activity[conv_id]
        logger.info(f"Cleared {len(inactive_convos)} inactive conversations")
        await asyncio.sleep(60)  # Check every minute

# ============================================================================
# FastAPI Events and Endpoints
# ============================================================================

@app.on_event("startup")
async def startup_event():
    logger.info("Starting up the application")
    init_db()
    asyncio.create_task(clear_inactive_conversations())

@app.post("/llm-agent")
async def llm_agent(query: LLMAgentQueryModel, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
    """
    LLM agent endpoint that provides responses based on user queries, maintaining conversation history.
    Accepts custom system messages and allows selection of different models.
    Requires API Key authentication via X-API-Key header.
    """
    logger.info(f"Received LLM agent query: {query.prompt}")

    # Generate a new conversation ID if not provided
    if not query.conversation_id:
        query.conversation_id = str(uuid4())

    # Initialize or retrieve conversation history
    if query.conversation_id not in conversations:
        system_message = query.system_message or "You are a helpful assistant. Provide concise and accurate responses."
        conversations[query.conversation_id] = [
            {"role": "system", "content": system_message}
        ]
    elif query.system_message:
        # Update system message if provided
        conversations[query.conversation_id][0] = {"role": "system", "content": query.system_message}

    # Add user's prompt to conversation history
    conversations[query.conversation_id].append({"role": "user", "content": query.prompt})
    last_activity[query.conversation_id] = time.time()

    # Limit tokens in the conversation history
    limited_conversation = limit_conversation_history(conversations[query.conversation_id])

    def process_response():
        full_response = ""
        for content in chat_with_llama_stream(limited_conversation, model=query.model_id):
            full_response += content
            yield content

        # Add the assistant's response to the conversation history
        conversations[query.conversation_id].append({"role": "assistant", "content": full_response})

        background_tasks.add_task(update_db, query.user_id, query.conversation_id, query.prompt, full_response)
        logger.info(f"Completed LLM agent response for query: {query.prompt}")

    return StreamingResponse(process_response(), media_type="text/event-stream")

import edge_tts
import io


@app.get("/tts")
async def text_to_speech(
    text: str = Query(..., description="Text to convert to speech"),
    voice: str = Query(default="en-GB-SoniaNeural", description="Voice to use for speech")
):
    communicate = edge_tts.Communicate(text, voice)

    async def generate():
        async for chunk in communicate.stream():
            if chunk["type"] == "audio":
                yield chunk["data"]

    return StreamingResponse(generate(), media_type="audio/mpeg")

# ============================================================================
# Main Execution
# ============================================================================

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)