import os
import json
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
import requests
import uvicorn

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

HF_TOKEN = os.getenv("HF_TOKEN")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
MEMORY_FILE = "memory.json"
MAIN_MODEL = "openai/gpt-oss-120b"

def load_memory():
    if os.path.exists(MEMORY_FILE):
        try:
            with open(MEMORY_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except: return []
    return []

def save_memory(user_msg, ai_res):
    memory = load_memory()
    memory.append({"role": "user", "content": user_msg})
    memory.append({"role": "assistant", "content": ai_res})
    if len(memory) > 20: memory = memory[-20:]
    with open(MEMORY_FILE, "w", encoding="utf-8") as f:
        json.dump(memory, f, ensure_ascii=False, indent=2)

@app.post("/chat")
async def chat(request: Request):
    try:
        user_data = await request.json()
        messages = user_data.get("messages", [])
        
        # 🧠 පරණ මතකයන් ලබා ගැනීම
        past_context = load_memory()
        
        # 🛡️ නිවැරදි කළ System Instructions (සිංහල අනිවාර්ය කර නැත)
        system_instructions = {
            "role": "system", 
            "content": "You are MINZO AI (v10.0), a strategic specialist. Respond in the same language the user uses. Be adaptive and evolve based on past memories."
        }
        
        full_payload_messages = [system_instructions] + past_context + messages

        # 🚀 GPT-OSS-120B හරහා දත්ත සැකසීම
        headers = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"}
        payload = {
            "model": MAIN_MODEL,
            "messages": full_payload_messages,
            "max_tokens": 1024,
            "temperature": 0.7
        }

        response = requests.post(
            "https://router.huggingface.co/openai/v1/chat/completions",
            headers=headers,
            json=payload,
            timeout=60
        )

        # Fallback to Groq
        if response.status_code != 200:
            fallback_res = requests.post(
                "https://api.groq.com/openai/v1/chat/completions",
                headers={"Authorization": f"Bearer {GROQ_API_KEY}"},
                json={"model": "llama-3.3-70b-versatile", "messages": full_payload_messages}
            )
            res_data = fallback_res.json()
        else:
            res_data = response.json()

        ai_reply = res_data["choices"][0]["message"]["content"]

        # Save to memory (if it's not a complex object)
        last_msg = messages[-1]["content"]
        if isinstance(last_msg, str):
            save_memory(last_msg, ai_reply)

        return res_data
        
    except Exception as e:
        return {"error": "Neural Core Failure", "details": str(e)}

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