import os
import re
import json
import uuid

from dotenv import load_dotenv
load_dotenv()

import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import google.genai as genai
import anthropic
import openai

app = FastAPI()

# In-memory store for chat sessions
chat_sessions = {}

# --- Utility Functions ---

def format_full_prompt(context, prompt):
    if isinstance(context, (dict, list)):
        formatted_context = "Context:\n" + json.dumps(context, indent=2)
    elif isinstance(context, str):
        formatted_context = "Context:\n" + context
    else:
        formatted_context = "Context:\n" + str(context)
    return f"{formatted_context}\n\nPrompt:\n{prompt}"

def clean_ai_response(raw_reply, general_call):
    raw_reply = raw_reply.strip()
    
    # Check for JSON block
    json_match = re.search(r'```json\s*(\{.*?\})\s*```', raw_reply, re.DOTALL)
    if json_match:
        try:
            return json.loads(json_match.group(1))
        except json.JSONDecodeError:
            pass

    if general_call != True:
        # Check for HTML block
        html_match = re.search(r'```html\s*(<table.*?</table>)\s*```', raw_reply, re.DOTALL)
        if html_match:
            return html_match.group(1)
        elif "<table" in raw_reply and "</table>" in raw_reply:
            table_match = re.search(r'(<table.*?</table>)', raw_reply, re.DOTALL)
            if table_match:
                return table_match.group(1)
    
    return raw_reply

# --- Vendor Logic Handlers ---

async def call_gemini(full_prompt):
    client = genai.Client(api_key=os.getenv('GEMINI_API_KEY'))
    response = client.models.generate_content(model=os.getenv("GEMINI_MODEL"), contents=full_prompt)
    return response.text

async def call_claude(full_prompt):
    client = anthropic.Anthropic(api_key=os.getenv('CLAUDE_API_KEY'))
    response = client.messages.create(
        model=os.getenv("CLAUDE_MODEL"),
        max_tokens=4096,
        messages=[{"role": "user", "content": full_prompt}]
    )
    return response.content[0].text

async def call_openai(full_prompt):
    client = openai.OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
    response = client.chat.completions.create(
        model=os.getenv("OPENAI_MODEL"),
        messages=[{"role": "user", "content": full_prompt}]
    )
    return response.choices[0].message.content

# --- Shared Logic ---

async def process_chat_request(data, vendor=None):
    if not vendor:
        vendor = data.get("vendor", "").lower()
    
    prompt = data.get("prompt")
    context = data.get("context", {})
    chat_session = data.get("chat_session")
    general_call = data.get("general_call", False)

    if not prompt or not chat_session:
        return JSONResponse({"error": "Missing 'prompt' or 'chat_session'"}, status_code=400)
    
    if not vendor:
        return JSONResponse({"error": "Missing 'vendor'"}, status_code=400)

    # Handle chat session
    if chat_session.lower() == "new":
        session_id = str(uuid.uuid4())
        chat_sessions[session_id] = {"type": vendor}
    else:
        session_id = chat_session

    # Validate context
    if general_call != True:
        if not isinstance(context, dict):
            return JSONResponse({"error": "Context must be a valid JSON object"}, status_code=400)

    full_prompt = format_full_prompt(context, prompt)

    try:
        if vendor == "gemini":
            raw_reply = await call_gemini(full_prompt)
        elif vendor == "claude":
            raw_reply = await call_claude(full_prompt)
        elif vendor == "openai":
            raw_reply = await call_openai(full_prompt)
        else:
            return JSONResponse({"error": f"Unsupported vendor: {vendor}"}, status_code=400)

        cleaned_reply = clean_ai_response(raw_reply, general_call)

        return JSONResponse({
            "session_id": session_id,
            "response": cleaned_reply,
            "vendor": vendor
        })
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)

# --- Routes ---

@app.post("/api/v0/chat")
async def unified_chat(request: Request):
    data = await request.json()
    return await process_chat_request(data)

@app.post("/api/v0/gemini-chat")
async def gemini_chat(request: Request):
    data = await request.json()
    return await process_chat_request(data, vendor="gemini")

@app.post("/api/v0/claude-chat")
async def claude_chat(request: Request):
    data = await request.json()
    return await process_chat_request(data, vendor="claude")

@app.post("/api/v0/openai-chat")
async def openai_chat(request: Request):
    data = await request.json()
    return await process_chat_request(data, vendor="openai")

if __name__ == "__main__":
    uvicorn.run("app:app", host="127.0.0.1", port=5007, reload=True, log_level="debug")