main.py

import os
from fastapi import FastAPI, Request
from fastapi.responses import PlainTextResponse, FileResponse
import requests

from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()

# Cấu hình CORS để cho phép Chrome Extension truy cập
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Đọc Java client code từ file bên ngoài
def get_java_client_code(filename: str):
    try:
        with open(filename, "r", encoding="utf-8") as f:
            return f.read()
    except Exception as e:
        return f"Error reading {filename}: {str(e)}"

@app.get("/get-java")
def get_java():
    return PlainTextResponse(content=get_java_client_code("Test2.java"))

@app.get("/get-java8")
def get_java8():
    return PlainTextResponse(content=get_java_client_code("Test3.java"))

@app.get("/get-c")
def get_c():
    return PlainTextResponse(content=get_java_client_code("client.c"))

@app.get("/get-c-curl")
def get_c_curl():
    return PlainTextResponse(content=get_java_client_code("client_curl.c"))

@app.get("/get-c-curl-pipe")
def get_c_curl_pipe():
    return PlainTextResponse(content=get_java_client_code("client_curl_pipe.c"))

@app.get("/get-cpp")
def get_cpp():
    return PlainTextResponse(content=get_java_client_code("client.cpp"))


@app.get("/download-extension")
def download_extension():
    extension_path = "chrome-extension.zip"
    if os.path.exists(extension_path):
        return FileResponse(
            path=extension_path,
            filename="chrome-extension.zip",
            media_type="application/zip"
        )
    return PlainTextResponse(content="Extension file not found", status_code=404)

# Lấy API Key từ Secret của Hugging Face
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")

@app.post("/ask")
async def ask(request: Request):
    try:
        # Nhận model id từ header hoặc dùng mặc định
        model_id = request.headers.get("X-Model", "poolside/laguna-s-2.1:free")
        
        # Thử parse body dưới dạng JSON trước
        try:
            data = await request.json()
            # Nếu là JSON, ta giả định nó có cấu trúc messages hoặc prompt + image_data
            if "messages" in data:
                messages = data["messages"]
            elif "image_data" in data:
                # Cấu trúc cho multi-modal 
                messages = [{
                    "role": "user",
                    "content": [
                        {"type": "text", "text": data.get("prompt", "")},
                        {
                            "type": "image_url",
                            "image_url": {"url": data["image_data"]}
                        }
                    ]
                }]
            else:
                messages = [{"role": "user", "content": data.get("prompt", "")}]
        except:
            # Nếu không phải JSON, dùng body thô làm prompt
            body = await request.body()
            prompt = body.decode('utf-8')
            messages = [{"role": "user", "content": prompt}]

        response = requests.post(
            url="https://openrouter.ai/api/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {OPENROUTER_API_KEY}",
                "Content-Type": "application/json",
            },
            json={
                "model": model_id, 
                "messages": messages
            }
        )
        
        resp_json = response.json()
        if 'choices' in resp_json:
            result = resp_json['choices'][0]['message']['content']
        else:
            result = f"OpenRouter Error: {str(resp_json)}"
        return PlainTextResponse(content=result)
    except Exception as e:
        return PlainTextResponse(content=f"Error: {str(e)}", status_code=500)

@app.get("/")
def read_root():
    return {"Status": "Proxy is running!"}