import os
import json
from typing import List, Dict, Any
from fastapi import FastAPI, HTTPException
from dotenv import load_dotenv
from supabase import create_client, Client
from groq import Groq
from google import genai
import uvicorn

load_dotenv()

app = FastAPI(title="Movie Linker API (Production)")

# --- Configuration ---
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_ANON_KEY")
GROQ_API_KEYS = [k.strip() for k in os.getenv("GROQ_API_KEYS", "").split(",") if k.strip()]
GEMINI_API_KEYS = [k.strip() for k in os.getenv("GEMINI_API_KEYS", "").split(",") if k.strip()]

if not SUPABASE_URL or not SUPABASE_KEY:
    raise ValueError("Missing Supabase credentials in environment variables.")

supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)

class APIKeyManager:
    """Manages rotation between Groq and Gemini providers and their keys."""
    def __init__(self, groq_keys, gemini_keys):
        self.keys = {"groq": groq_keys, "gemini": gemini_keys}
        self.indices = {"groq": 0, "gemini": 0}
        self.provider = "groq" if groq_keys else "gemini"

    def get_client(self):
        if not self.keys[self.provider]:
            if self.provider == "groq" and self.keys["gemini"]:
                self.provider = "gemini"
                return self.get_client()
            raise Exception("No API keys available.")
        
        current_key = self.keys[self.provider][self.indices[self.provider]]
        if self.provider == "groq":
            return Groq(api_key=current_key), "groq"
        return genai.Client(api_key=current_key), "gemini"

    def rotate(self):
        if self.indices[self.provider] + 1 < len(self.keys[self.provider]):
            self.indices[self.provider] += 1
            print(f"🔄 Switched to {self.provider} key #{self.indices[self.provider]+1}")
        elif self.provider == "groq" and self.keys["gemini"]:
            self.provider = "gemini"
            print("⚠️ Groq exhausted, falling back to Gemini.")
        else:
            raise Exception("All API keys for all providers exhausted.")

key_manager = APIKeyManager(GROQ_API_KEYS, GEMINI_API_KEYS)

# --- Core Functions ---

def ai_match(client, provider, prompt):
    if provider == "groq":
        res = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model="llama-3.3-70b-versatile",
            response_format={"type": "json_object"}
        )
        return json.loads(res.choices[0].message.content)
    
    res = client.models.generate_content(
        model="gemini-3-flash-preview",
        contents=prompt,
        config={'response_mime_type': 'application/json'}
    )
    return json.loads(res.text)

def get_db_data(table: str, select: str = "*"):
    return supabase.table(table).select(select).execute().data

def process_batch(shorts_batch, movies_ref):
    while True:
        try:
            client, provider = key_manager.get_client()
            prompt = f"Match these shorts to media. Media: {json.dumps(movies_ref[:120])}. Shorts: {json.dumps(shorts_batch)}. Respond with JSON: {{'matches': [{{'short_id', 'media_id', 'short_title', 'media_title', 'short_link', 'media_link'}}]}}"
            return ai_match(client, provider, prompt).get("matches", [])
        except Exception as e:
            if any(x in str(e).lower() for x in ["rate_limit", "429", "limit_reached"]):
                key_manager.rotate()
            else:
                print(f"❌ Batch Error: {e}")
                return []

# --- Endpoints ---

@app.get("/")
def home():
    return {"status": "online", "engine": "Groq+Gemini"}

@app.get("/status")
def status():
    linked = get_db_data("linked_results", "short_id")
    return {
        "linked_count": len(linked),
        "provider": key_manager.provider,
        "active_keys": {p: len(k) for p, k in key_manager.keys.items()}
    }

@app.post("/link-data")
def link_data():
    shorts = get_db_data("shorts")
    media = get_db_data("media")
    linked_ids = {str(l["short_id"]) for l in get_db_data("linked_results", "short_id")}
    
    # Filter only new shorts
    new_shorts = [s for s in shorts if str(s.get("id")) not in linked_ids]
    if not new_shorts: return {"message": "Everything is already linked."}

    # Prepare data for AI
    movies_ref = [{
        "id": m["id"], "title": m["title"], "type": m.get("type"),
        "year": str(m.get("releasdate", ""))[:4], "desc": m.get("dec", "")[:100]
    } for m in media]

    batch_shorts = [{
        "id": s["id"], "title": s.get("title"), "desc": s.get("description", "")[:100]
    } for s in new_shorts]

    # Matching Loop
    all_matches = []
    for i in range(0, len(batch_shorts), 20):
        all_matches.extend(process_batch(batch_shorts[i : i + 20], movies_ref))

    # Save Results
    if all_matches:
        supabase.table("linked_results").insert(all_matches).execute()
    
    return {"newly_linked": len(all_matches), "total": len(linked_ids) + len(all_matches)}

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