import json
import asyncio
from datetime import datetime, timedelta
from typing import Type, AsyncGenerator, List

from bson import json_util
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from pymongo import MongoClient

from langchain_community.llms import Ollama
from langchain.agents import initialize_agent, AgentType
from langchain.schema import HumanMessage, AIMessage, SystemMessage
from langchain.prompts.chat import ChatPromptTemplate
from langchain.tools import BaseTool, tool
from langchain.callbacks.base import AsyncCallbackHandler

from sse_starlette.sse import EventSourceResponse

# =========================
# MongoDB
# =========================
MONGO_URI = "mongodb://fikoyt:fikoyt1299@fikoyt.com:9001/?replicaSet=rs0&tls=false&authSource=fikoyt"
DB_NAME = "fikoyt"
COLLECTION_USERS = "goUyeler"
COLLECTION_CHAT = "fikoGPT_chat_history"

client = MongoClient(MONGO_URI)
db = client[DB_NAME]
chat_col = db[COLLECTION_CHAT]

# =========================
# LLM Ayarları
# =========================
BASE_LLM_KW = dict(
    model="gpt-oss",   # Kullanıcının verdiği isim
    stream=True,
    options={
        "num_predict": 512,
        "temperature": 0.4,
    },
)

system_prompt_text = """
Sen FikoGPT adında, Türkçe konuşan bağlama duyarlı bir yapay zekâsın.
Kısa, net ve doğru cevaplar ver. Özel verileri sorma veya ifşa etme.
"""

chat_prompt = ChatPromptTemplate.from_messages([
    SystemMessage(content=system_prompt_text),
    HumanMessage(content="{user_input}")
])

# =========================
# Araçlar
# =========================
class HelloToolInput(BaseModel):
    name: str = Field(description="Kullanıcıya hitap edilecek isim")

class MerhabaTool(BaseTool):
    name: str = "merhaba_ver"
    description: str = "Belirtilen isme Merhaba diyerek selam verir."
    args_schema: Type[BaseModel] = HelloToolInput
    def _run(self, name: str) -> str:
        return f"Merhaba, {name}!"
    def _arun(self, name: str):
        raise NotImplementedError()

@tool
def tarih_ver(asd=None) -> str:
    """
    Kullanıcıya güncel tarih ve saati verir.
    - Gün, ay, yıl, saat, dakika, saniye bilgilerini içerir.
    - Dün, bugün, yarın ve ertesi gün bilgilerini içerir.
    - Geçen ay, bu ay ve gelecek ay bilgilerini içerir.
    - Yılın kaçıncı haftasında olduğunu da içerir.
    """
    ay_adlari = {
        1: "Ocak", 2: "Şubat", 3: "Mart", 4: "Nisan",
        5: "Mayıs", 6: "Haziran", 7: "Temmuz", 8: "Ağustos",
        9: "Eylül", 10: "Ekim", 11: "Kasım", 12: "Aralık"
    }
    gun_adlari = {
        0: "Pazartesi", 1: "Salı", 2: "Çarşamba", 3: "Perşembe",
        4: "Cuma", 5: "Cumartesi", 6: "Pazar"
    }
    simdi = datetime.now()
    def gun_bilgisi(tarih):
        return {
            "tarih": tarih.strftime("%d.%m.%Y %H:%M:%S"),
            "gun": gun_adlari[tarih.weekday()],
            "ay": ay_adlari[tarih.month],
            "yil": tarih.year,
            "hafta": int(tarih.strftime("%W")) + 1,
        }
    dun = simdi - timedelta(days=1)
    yarin = simdi + timedelta(days=1)
    ertesi = simdi + timedelta(days=2)
    return json.dumps({
        "bugun": gun_bilgisi(simdi),
        "dun": gun_bilgisi(dun),
        "yarin": gun_bilgisi(yarin),
        "ertesi": gun_bilgisi(ertesi),
        "bu_ay": {"ay": ay_adlari[simdi.month], "yil": simdi.year},
    }, ensure_ascii=False)

class MongoDBToolInput(BaseModel):
    name: str = Field(
        default="all",
        description="MongoDB sorgusu: 'all', 'alan:değer' veya sadece isim"
    )

class MongoDBTool(BaseTool):
    name: str = "mongodb_kullanici_sorgula"
    description: str = "MongoDB'den kullanıcı sorgular."
    args_schema: Type[BaseModel] = MongoDBToolInput
    def _run(self, name: str = "all") -> str:
        collection = db[COLLECTION_USERS]
        if not name or name.lower() in ["all", "hepsi", "tamamı"]:
            users = list(collection.find({}))
            return json.dumps(users, ensure_ascii=False, default=json_util.default) if users else "Hiç kullanıcı yok."
        if ":" in name:
            field, value = name.split(":", 1)
            query = {field.strip(): {"$in": [value.strip()]}}
            users = list(collection.find(query))
            return json.dumps(users, ensure_ascii=False, default=json_util.default) if users else f"{field}={value} bulunamadı."
        user = collection.find_one({"isim": name})
        return json.dumps(user, ensure_ascii=False, default=json_util.default) if user else f"{name} bulunamadı."
    def _arun(self, name: str = "all"):
        raise NotImplementedError()

tools = [MongoDBTool(), MerhabaTool(), tarih_ver]

# =========================
# Agent
# =========================
def build_llm(callbacks=None):
    return Ollama(callbacks=callbacks, **BASE_LLM_KW)

def build_agent(callbacks=None):
    return initialize_agent(
        tools,
        build_llm(callbacks),
        agent=AgentType.OPENAI_FUNCTIONS,
        max_iterations=3,
        verbose=False,
        handle_parsing_errors=True,
    )

# =========================
# FastAPI
# =========================
app = FastAPI(title="FikoGPT API", version="1.2.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Static files (frontend)
app.mount("/static", StaticFiles(directory="static"), name="static")

class AskRequest(BaseModel):
    question: str

def load_chat_history(limit: int = 10) -> List:
    msgs = []
    for doc in chat_col.find().sort("timestamp", -1).limit(limit):
        if doc["role"] == "human":
            msgs.append(HumanMessage(content=doc["message"]))
        elif doc["role"] == "ai":
            msgs.append(AIMessage(content=doc["message"]))
    msgs.reverse()
    return msgs

@app.post("/ask")
def ask(req: AskRequest):
    chat_history = load_chat_history(10)
    agent = build_agent()
    answer = agent.run(input=req.question, chat_history=chat_history)
    chat_col.insert_many([
        {"role": "human", "message": req.question, "timestamp": datetime.utcnow()},
        {"role": "ai", "message": answer, "timestamp": datetime.utcnow()}
    ])
    return {"answer": answer}

# =========================
# Streaming (SSE)
# =========================
class QueueStreamCallback(AsyncCallbackHandler):
    """Tokenları asyncio.Queue'a yazar, SSE/WS tüketir."""
    def __init__(self, queue: asyncio.Queue):
        self.queue = queue

    async def on_llm_new_token(self, token: str, **kwargs) -> None:
        await self.queue.put(token)

    async def on_llm_end(self, *args, **kwargs) -> None:
        # Bitti sinyali
        await self.queue.put("[[END]]")

async def run_agent_stream(question: str, queue: asyncio.Queue):
    """Agent'ı başlatır ve bitince DB'ye yazar."""
    agent = build_agent(callbacks=[QueueStreamCallback(queue)])
    chat_history = load_chat_history(10)
    answer = await agent.arun(input=question, chat_history=chat_history)
    chat_col.insert_many([
        {"role": "human", "message": question, "timestamp": datetime.utcnow()},
        {"role": "ai", "message": answer, "timestamp": datetime.utcnow()}
    ])

async def sse_generator(question: str) -> AsyncGenerator[str, None]:
    queue: asyncio.Queue = asyncio.Queue()
    task = asyncio.create_task(run_agent_stream(question, queue))
    try:
        while True:
            token = await queue.get()
            if token == "[[END]]":
                break
            yield token
        await task
    except Exception as e:
        yield f"[HATA] {str(e)}"

@app.get("/ask/stream")
async def ask_stream(question: str):
    return EventSourceResponse(sse_generator(question), ping=15000)

# =========================
# WebSocket
# =========================
@app.websocket("/ws")
async def ws_endpoint(ws: WebSocket):
    await ws.accept()
    try:
        while True:
            question = await ws.receive_text()
            queue: asyncio.Queue = asyncio.Queue()
            task = asyncio.create_task(run_agent_stream(question, queue))
            try:
                while True:
                    token = await queue.get()
                    if token == "[[END]]":
                        await ws.send_text("[END]")
                        break
                    await ws.send_text(token)
                await task
            except Exception as e:
                await ws.send_text(f"[HATA] {str(e)}")
    except WebSocketDisconnect:
        pass
    except Exception as e:
        try:
            await ws.send_text(f"[HATA] {str(e)}")
        except:
            pass
        await ws.close()

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