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

app = FastAPI()

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

ModelID = Literal[
    "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"
]

class QueryModel(BaseModel):
    user_query: str = Field(..., description="User's coding query")
    model_id: ModelID = Field(
        default="meta-llama/llama-3-70b-instruct",
        description="ID of the model to use for response generation"
    )
    conversation_id: str = Field(default_factory=lambda: str(uuid4()), description="Unique identifier for the conversation")
    user_id: str = Field(..., description="Unique identifier for the user")

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

@lru_cache()
def get_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")

def limit_tokens(input_string, token_limit=6000):
    return encoding.decode(encoding.encode(input_string)[:token_limit])

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

async def chat_with_llama_stream(messages, model="gpt-3.5-turbo", max_llm_history=4, max_output_tokens=2500):
    while calculate_tokens(messages) > (8000 - max_output_tokens):
        if len(messages) > max_llm_history:
            messages = [messages[0]] + messages[-max_llm_history:]
        else:
            max_llm_history -= 1
            if max_llm_history < 2:
                raise ValueError("Unable to reduce message length below token limit")

    try:
        response = await or_client.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=max_output_tokens,
            stream=True
        )
        
        full_response = ""
        async 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})
        return full_response
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error in model response: {str(e)}")


async def verify_api_key(api_key: str = Security(api_key_header)):
    if api_key != API_KEY:
        raise HTTPException(status_code=403, detail="Could not validate credentials")
    return api_key

# SQLite setup
DB_PATH = '/app/data/conversations.db'

def init_db():
    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()

init_db()

def update_db(user_id, conversation_id, message, response):
    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()

async def clear_inactive_conversations():
    while True:
        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]
        await asyncio.sleep(60)  # Check every minute

@app.on_event("startup")
async def startup_event():
    asyncio.create_task(clear_inactive_conversations())

@app.post("/coding-assistant")
async def coding_assistant(query: QueryModel, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
    """
    Coding assistant endpoint that provides programming help based on user queries.

    Available models:
    - meta-llama/llama-3-70b-instruct (default)
    - 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

    Requires API Key authentication via X-API-Key header.
    """
    if query.conversation_id not in conversations:
        conversations[query.conversation_id] = [
            {"role": "system", "content": "You are a helpful assistant proficient in coding tasks. Help the user in understanding and writing code."}
        ]
    
    conversations[query.conversation_id].append({"role": "user", "content": query.user_query})
    last_activity[query.conversation_id] = time.time()
    
    # Limit tokens in the conversation history
    limited_conversation = conversations[query.conversation_id]
    while calculate_tokens(limited_conversation) > 8000:
        if len(limited_conversation) > 2:  # Keep at least the system message and the latest user message
            limited_conversation.pop(1)
        else:
            error_message = "Token limit exceeded. Please shorten your input or start a new conversation."
            raise HTTPException(status_code=400, detail=error_message)

    async def process_response():
        full_response = ""
        async for content in chat_with_llama_stream(limited_conversation, model=query.model_id):
            full_response += content
            yield content
        background_tasks.add_task(update_db, query.user_id, query.conversation_id, query.user_query, full_response)

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

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