from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import requests
import json
import re
import time

app = FastAPI()

class ChatRequest(BaseModel):
    cookies: str
    prompt: str
    chat_id: str = None

class GeminiAPI:
    def __init__(self, cookies_string):
        self.cookies = self._parse_cookies(cookies_string)
        self.session = requests.Session()
        self.session.cookies.update(self.cookies)
        
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (Linux; Android 10; K) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Mobile Safari/537.36',
            'Content-Type': 'application/json+protobuf',
            'Accept': '*/*',
            'Origin': 'https://aistudio.google.com',
            'Referer': 'https://aistudio.google.com/',
            'X-Goog-Api-Key': 'AIzaSyD0fP816MBRBSkj20Q04XbjsIcI0jGI0WOs',
            'X-Client-Data': 'CJGVywE=',
        }
        
    def _parse_cookies(self, cookies_input):
        cookies = {}
        
        try:
            cookie_list = json.loads(cookies_input)
            if isinstance(cookie_list, list):
                for cookie in cookie_list:
                    if 'name' in cookie and 'value' in cookie:
                        cookies[cookie['name']] = cookie['value']
                return cookies
        except:
            pass
        
        if isinstance(cookies_input, str):
            parts = cookies_input.split(';')
            for part in parts:
                if '=' in part:
                    name, value = part.strip().split('=', 1)
                    cookies[name] = value
        
        return cookies
    
    def send_message(self, prompt, chat_id=None):
        url = 'https://alkalimakersuite-pa.clients6.google.com/$rpc/google.internal.alkali.applications.makersuite.v1.MakerSuiteService/GenerateContent'
        
        payload = [
            "models/gemini-3-flash-preview",
            [[[[None, prompt]], "user"]],
            None,
            0,
            "models/gemini-3-flash-preview",
            [None, None, None, 65536, 1, 0.95, 64, None, None, None, None, None, None, None],
            "!GBu1G0PNAAZ19fabc_VClWQf4_wu6go7ADQBEArZ1NrXSqgcKEzmXhn1r70At0ogWMU1G6bOhL3",
            None,
            [[None, None, [[]]],
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            [[None, None, "Asia/Makassar"]]
        ]
        
        try:
            response = self.session.post(
                url,
                headers=self.headers,
                json=payload,
                timeout=30
            )
            
            if response.status_code != 200:
                return None, f"HTTP Error: {response.status_code}"
            
            data = response.json()
            
            text_response = ""
            
            if isinstance(data, list) and len(data) > 0:
                for item in data:
                    if isinstance(item, list) and len(item) > 0:
                        for subitem in item:
                            if isinstance(subitem, list) and len(subitem) > 0:
                                for content in subitem:
                                    if isinstance(content, list) and len(content) > 0:
                                        for text_item in content:
                                            if isinstance(text_item, list) and len(text_item) > 0:
                                                if text_item[0] is None and len(text_item) > 1:
                                                    if isinstance(text_item[1], str):
                                                        text_response += text_item[1]
            
            if text_response:
                return {
                    'response': text_response.strip(),
                    'chat_id': chat_id,
                    'timestamp': int(time.time())
                }, None
            
            return None, "Tidak bisa parse response"
            
        except Exception as e:
            return None, str(e)

@app.post("/chat")
async def chat(request: ChatRequest):
    try:
        gemini = GeminiAPI(request.cookies)
        result, error = gemini.send_message(request.prompt, request.chat_id)
        
        if error:
            raise HTTPException(status_code=400, detail=error)
        
        return result
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return {
        "message": "Gemini API",
        "version": "2.0",
        "endpoints": {
            "POST /chat": {
                "cookies": "string (JSON array atau cookie string)",
                "prompt": "string",
                "chat_id": "string (optional)"
            }
        }
    }

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