import os
from fastapi import FastAPI
from pydantic import BaseModel
from huggingface_hub import InferenceClient
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse

app = FastAPI()

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

# Connect to the stable Qwen-7B Brain (100% Free on HF Cloud)
client = InferenceClient(
    model="Qwen/Qwen2.5-7B-Instruct",
    token=os.getenv("HF_TOKEN")
)

memory_store = {}

# CUSTOMIZABLE IDENTITY: Your friend can edit this text!
AI_IDENTITY = """
You are a highly advanced, warm, and supportive AI companion.
Your name is ZENVY AI
You were built by Rik Sarkar
Always be encouraging, speak naturally, and help the user.
"""

class ChatRequest(BaseModel):
    message: str
    session_id: str
  
@app.get("/")
def home():
    return {"status": "ZENVY V1 is Online"}

@app.post("/chat")
async def chat(request: ChatRequest):
    session_id = request.session_id
    if session_id not in memory_store:
        memory_store[session_id] = []
        
    history = memory_store[session_id]
    history.append({"role": "user", "content": request.message})

    def event_generator():
        collected_reply = []
        try:
            messages = [{"role": "system", "content": AI_IDENTITY}] + history
            
            response = client.chat_completion(
                messages=messages,
                max_tokens=4096,  # Raised output token capacity for long, complete responses
                temperature=0.7,
                stream=True
            )
            
            for chunk in response:
                if chunk.choices and len(chunk.choices) > 0:
                    word = chunk.choices[0].delta.content
                    if word:
                        collected_reply.append(word)
                        yield word
            
            full_reply = "".join(collected_reply)
            history.append({"role": "assistant", "content": full_reply})
            
            if len(history) > 12:
                history.pop(0)
                history.pop(0)
                
        except Exception as e:
            yield f"Error: {str(e)}"

    return StreamingResponse(event_generator(), media_type="text/event-stream")