import os
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import List, Dict
import google.generativeai as genai
from google.generativeai.types import HarmBlockThreshold, HarmCategory
from dotenv import load_dotenv

load_dotenv()

app = FastAPI()

class ChatMessage(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[ChatMessage]
    api_key: str = None

chat_history = []

async def stream_response(response):
    for chunk in response:
        yield chunk.text

@app.post("/api/think")
async def think(request: ChatRequest):
    api_key = request.api_key or os.getenv("GOOGLE_API_KEY")
    if not api_key:
        raise HTTPException(status_code=400, detail="API key is required")

    genai.configure(api_key=api_key)
    model = genai.GenerativeModel(
        model_name="gemini-1.5-flash",
        generation_config={
            "temperature": 1,
            "max_output_tokens": 8192,
            "response_mime_type": "application/json",
        },
        safety_settings={
            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE
        },
        system_instruction="You are a helpful AI assistant. Analyze the user's request and provide your thoughts in XML tags. Use <thinking> for overall analysis, <step> for individual steps, and <count> for remaining steps."
    )

    messages = [{"role": m.role, "parts": [m.content]} for m in request.messages]
    response = await model.generate_content_async(messages, stream=True)
    return StreamingResponse(stream_response(response), media_type="text/plain")

@app.post("/api/chat")
async def chat(request: ChatRequest):
    api_key = request.api_key or os.getenv("GOOGLE_API_KEY")
    if not api_key:
        raise HTTPException(status_code=400, detail="API key is required")

    genai.configure(api_key=api_key)
    model = genai.GenerativeModel(
        model_name="gemini-1.5-flash",
        generation_config={
            "temperature": 1,
            "max_output_tokens": 8192,
            "response_mime_type": "application/json",
        },
        safety_settings={
            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE
        },
        system_instruction="You are a helpful AI assistant. Provide a concise summary of the thinking process and give a final answer."
    )

    messages = [{"role": m.role, "parts": [m.content]} for m in request.messages]
    response = await model.generate_content_async(messages, stream=True)
    return StreamingResponse(stream_response(response), media_type="text/plain")

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