import os
import glob
import json
import ast
import pandas as pd
import numpy as np
import requests
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import uvicorn
from huggingface_hub import snapshot_download
from llama_cpp import Llama

app = FastAPI()

# --- 1. CONFIG & DATA ---
BASE_DIR = os.path.dirname(__file__)
CSV_PATH = os.path.join(BASE_DIR, "Races Database_rows-embed.csv")
ACCOUNT_ID = os.environ.get("ACCOUNT_ID")
API_TOKEN = os.environ.get("API_TOKEN")

@app.on_event("startup")
def load_engine():
    global df_display, db_matrix, llm
    print("🚀 Initializing NextLap Engine...")
    
    # Load and Parse CSV (optimized for RAM)
    df = pd.read_csv(CSV_PATH).fillna("N/A")
    df['vec'] = df['embedding'].apply(lambda x: np.array(ast.literal_eval(x), dtype=np.float32))
    db_matrix = np.vstack(df['vec'].values)
    # Context Compression: Keep only vital columns for the LLM
    vital_cols = ['event', 'city', 'date', 'registrationCost', 'type']
    df_display = df[vital_cols]

    # Load Gemma 4 (n_threads=2 for HF Free Tier)
    model_dir = snapshot_download(repo_id="lmstudio-community/gemma-4-E2B-it-GGUF", allow_patterns=["*Q4_K_M.gguf"])
    model_path = glob.glob(f"{model_dir}/*Q4_K_M.gguf")[0]
    llm = Llama(model_path=model_path, n_ctx=2048, n_threads=2, n_batch=512)

# --- 2. INTENT ROUTER (SPEED OPTIMIZATION) ---
def should_trigger_db(query):
    """Determines if query needs DB. Prevents unnecessary API calls & latency."""
    keywords = [
        'race', 'marathon', 'triathlon', 'run', 'event', 'cost', 'price', 
        'date', 'when', 'mumbai', 'bangalore', 'goa', 'delhi', 'pune', 'india',
        'cycling', 'swimming', 'upcoming', 'fee', 'register'
    ]
    return any(k in query.lower() for k in keywords)

def embed_query(text):
    url = f"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/run/@cf/baai/bge-m3"
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    try:
        res = requests.post(url, headers=headers, json={"text": [text]}, timeout=4)
        return np.array(res.json()["result"]["data"][0], dtype=np.float32)
    except: return None

def get_context(query_text):
    if not should_trigger_db(query_text):
        return None # Chat normally
    
    query_vec = embed_query(query_text)
    if query_vec is None: return None
    
    # Fast Cosine Similarity
    sims = np.dot(db_matrix, query_vec) / (np.linalg.norm(db_matrix, axis=1) * np.linalg.norm(query_vec))
    
    # Top-K Reduction: Only 3 matches for maximum speed and relevance
    top_indices = np.argsort(sims)[-3:][::-1]
    results = df_display.iloc[top_indices]
    
    # Context Compression: Markdown table is the most token-efficient format
    return results.to_markdown(index=False)

# --- 3. STREAMING API ---
@app.post("/v1/chat/completions")
async def chat(request: Request):
    data = await request.json()
    messages = data.get("messages", [])
    user_query = messages[-1]['content']
    
    context = get_context(user_query)
    
    if context:
        # No Hallucination Guardrail: Strict instruction when context is present
        sys_prompt = (
"""
You are 'NextLap AI', a sports travel companion for athletes in India.
STRICT GUARDRAILS:
1. BE NATURAL: Be highly encouraging. If they mention doing their 'first' race, congratulate them.
2. NO HALLUCINATIONS: Only list races provided in the DATABASE RESULTS. If empty, clearly state you have no races for that criteria right now.
3. FORMATTING: Use bullet points or short tables. DO NOT invent links.

Users often ask tricky conversational questions (e.g., "I want to do my first triathlon this year").
Even if it sounds like general chat, IF they mention a sport, location, or timeframe, you MUST extract it!
Normalize `sport_type` to general categories: "running", "cycling", "swimming", "triathlon", or null.

"""
            f"DB:\n{context}"
        )
    else:
        # Normal LLM Behavior: Friendly persona for general chat
        sys_prompt = "You are NextLap AI, a friendly sports companion. Chat naturally."

    messages.insert(0, {"role": "system", "content": sys_prompt})

    def stream_generator():
        response_iter = llm.create_chat_completion(
            messages=messages,
            max_tokens=400,
            temperature=0.2,
            stream=True 
        )
        for chunk in response_iter:
            delta = chunk['choices'][0]['delta']
            if 'content' in delta:
                yield f"data: {json.dumps({'choices': [{'delta': {'content': delta['content']}}]})}\n\n"
        yield "data: [DONE]\n\n"

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

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