import asyncio
from fastapi import FastAPI, Request, WebSocket
from fastapi.responses import StreamingResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from langchain.agents import initialize_agent, AgentType
from langchain_community.llms import Ollama
from langchain.schema import HumanMessage, AIMessage, SystemMessage
from bson import json_util
from pymongo import MongoClient
from datetime import datetime

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

# ---------- MongoDB ----------
MONGO_URI = "mongodb://fikoyt:fikoyt1299@fikoyt.com:9001/?replicaSet=rs0&tls=false&authSource=fikoyt"
DB_NAME = "fikoyt"
COLLECTION_CHAT = "fikoGPT_chat_history"
client = MongoClient(MONGO_URI)
db = client[DB_NAME]
chat_col = db[COLLECTION_CHAT]

# ---------- LLM & Agent ----------
system_prompt_text = """
Sen FikoGPT'sin, tamamen Türkçe konuşursun.
Cevap verirken format: Thought, Action, Action Input
"""
chat_history = []

llm = Ollama(model="gpt-oss")  # callbacks ile SSE/WS ekleyebilirsin

tools = []  # İstersen burada MerhabaTool, tarih_ver vb ekleyebilirsin

agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
    max_iterations=5,
    verbose=True,
)

# ---------- SSE akışı ----------
class QueueStreamCallback:
    def __init__(self, queue: asyncio.Queue):
        self.queue = queue
    def on_llm_new_token(self, token: str, **kwargs):
        asyncio.create_task(self.queue.put(token))

async def run_agent_stream(question: str):
    queue = asyncio.Queue()
    llm_with_cb = Ollama(model="gpt-oss", callbacks=[QueueStreamCallback(queue)])
    # Burada agent.run input verilecek
    answer = agent.run(input=question, chat_history=chat_history)
    yield f"data: {answer}\n\n"

@app.get("/ask/stream")
async def ask_stream(question: str):
    return StreamingResponse(run_agent_stream(question), media_type="text/event-stream")

# ---------- WebSocket ----------
@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    await ws.accept()
    try:
        while True:
            data = await ws.receive_text()
            answer = agent.run(input=data, chat_history=chat_history)
            await ws.send_text(answer)
    except Exception as e:
        await ws.send_text(f"Hata: {e}")

# ---------- Frontend ----------
@app.get("/static/index.html")
async def get_index():
    with open("static/index.html", "r", encoding="utf-8") as f:
        return HTMLResponse(f.read())