import os
import uuid
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Optional, Dict

from langchain_groq import ChatGroq
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_message_histories import ChatMessageHistory

# ── App Setup 
app = FastAPI(
    title="NeuroBot API",
    description="Advanced Brain Tumor AI Assistant powered by LangChain + Groq",
    version="2.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
)

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

# ── Load system prompt 
with open("system_prompt.txt", "r", encoding="utf-8") as f:
    SYSTEM_PROMPT = f.read()

# ── LangChain LLM 
llm = ChatGroq(
    model="llama-3.3-70b-versatile",
    temperature=0.4,
    max_tokens=1024,
    api_key=os.getenv("GROQ_API_KEY"),
)

# ── In-memory session store 
session_store: Dict[str, ChatMessageHistory] = {}
MAX_SESSIONS = 500
MAX_HISTORY_MESSAGES = 30


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in session_store:
        if len(session_store) >= MAX_SESSIONS:
            oldest = next(iter(session_store))
            del session_store[oldest]
        session_store[session_id] = ChatMessageHistory()
    return session_store[session_id]


# ── LangChain chain with history 
prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}"),
])

chain = prompt | llm

chain_with_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

# ── Pydantic Models 
class ChatRequest(BaseModel):
    message: str = Field(..., min_length=1, max_length=2000)
    session_id: Optional[str] = Field(default=None)

class ChatResponse(BaseModel):
    reply: str
    session_id: str
    turn_count: int

class NewSessionResponse(BaseModel):
    session_id: str
    message: str

class SessionInfoResponse(BaseModel):
    session_id: str
    turn_count: int
    exists: bool

class HealthResponse(BaseModel):
    status: str
    model: str
    active_sessions: int
    version: str


# ── Endpoints 

@app.post("/chat", response_model=ChatResponse, summary="Send a message to NeuroBot")
async def chat(request: ChatRequest):
    session_id = request.session_id or str(uuid.uuid4())
    user_message = request.message.strip()

    history_obj = get_session_history(session_id)
    if len(history_obj.messages) > MAX_HISTORY_MESSAGES:
        history_obj.messages = history_obj.messages[-MAX_HISTORY_MESSAGES:]

    try:
        response = chain_with_history.invoke(
            {"input": user_message},
            config={"configurable": {"session_id": session_id}},
        )
        reply_text = response.content
        turn_count = len(get_session_history(session_id).messages)

        return ChatResponse(
            reply=reply_text,
            session_id=session_id,
            turn_count=turn_count,
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"LLM error: {str(e)}")


@app.post("/session/new", response_model=NewSessionResponse, summary="Create a new session")
async def new_session():
    session_id = str(uuid.uuid4())
    get_session_history(session_id)
    return NewSessionResponse(
        session_id=session_id,
        message="New session created. Use this session_id in your /chat requests."
    )


@app.get("/session/{session_id}", response_model=SessionInfoResponse, summary="Get session info")
async def session_info(session_id: str):
    exists = session_id in session_store
    turn_count = len(session_store[session_id].messages) if exists else 0
    return SessionInfoResponse(session_id=session_id, turn_count=turn_count, exists=exists)


@app.delete("/session/{session_id}", summary="Clear session history")
async def clear_session(session_id: str):
    if session_id in session_store:
        del session_store[session_id]
        return {"message": f"Session {session_id} cleared successfully."}
    raise HTTPException(status_code=404, detail="Session not found.")


@app.get("/", response_model=HealthResponse, summary="Health check")
async def health():
    return HealthResponse(
        status="NeuroBot is running successfully",
        model="llama-3.3-70b-versatile via LangChain + Groq",
        active_sessions=len(session_store),
        version="2.0.0",
    )