from fastapi import APIRouter, HTTPException, Request
from api.models import PromptRequest, HistoryRequest
from core.llm import llm, GROQ_MODEL
from core.memory import get_memory, conversation_memories
from core.utils import classify_message, process_response
import uuid
import time
from langchain.schema import HumanMessage, AIMessage
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain, ConversationChain

router = APIRouter()

# Prompt templates
CONVERSATION_TEMPLATE = """You are a helpful AI assistant. Have a natural conversation with the user.

Current conversation:
{history}
Human: {input}
AI:"""

CODE_GENERATION_TEMPLATE = """You are an expert Python programmer. Generate clean, efficient Python code for the given request.

Request: {input}

Provide only the Python code without explanation:
```python"""

EXPLANATION_TEMPLATE = """You are a programming tutor. Explain how to solve the programming task clearly without providing code.

Task: {input}

Explanation:"""

BOTH_TEMPLATE = """You are an expert Python programmer and tutor. For the given request, provide both a clear explanation and Python code.

Request: {input}

First, provide a clear explanation of the approach:
[EXPLANATION]

Then, provide the Python code:
```python"""

def create_chain(template: str, memory=None):
    input_vars = ["input"]
    if memory:
        input_vars.append("history")
    prompt = PromptTemplate(
        input_variables=input_vars,
        template=template
    )
    if memory:
        return ConversationChain(
            llm=llm,
            prompt=prompt,
            memory=memory,
            verbose=True,
            output_key="output"
        )
    else:
        return LLMChain(
            llm=llm,
            prompt=prompt,
            verbose=True,
            output_key="text"
        )

@router.get("/")
async def root():
    return {
        "status": "ok",
        "title": "LangChain Chat Bot",
        "version": "2.0.0",
        "model": GROQ_MODEL,
        "description": "Advanced chatbot using LangChain and Groq models",
        "features": ["Conversation Memory", "Code Generation", "Explanations", "Session Management"]
    }

@router.post("/generate/")
async def generate_response(request: PromptRequest):
    if llm is None:
        raise HTTPException(status_code=503, detail="AI model not initialized. Check server logs for errors during startup (e.g., missing API key).")
    try:
        session_id = request.session_id or str(uuid.uuid4())
        memory = get_memory(session_id, request.max_history)
        message_type = classify_message(request.prompt)
        response_data = {}
        if message_type == "conversation" or request.response_type == "conversation":
            chain = create_chain(CONVERSATION_TEMPLATE, memory)
            response = chain.predict(input=request.prompt)
            response_data = {"response": response.strip(), "message_type": "conversation"}
        elif request.response_type == "code":
            chain = create_chain(CODE_GENERATION_TEMPLATE)
            response = chain.run(input=request.prompt)
            response_data = process_response(response, "code")
            response_data["message_type"] = "code"
        elif request.response_type == "explanation":
            chain = create_chain(EXPLANATION_TEMPLATE)
            response = chain.run(input=request.prompt)
            response_data = process_response(response, "explanation")
            response_data["message_type"] = "explanation"
        else:  # "both"
            chain = create_chain(BOTH_TEMPLATE)
            response = chain.run(input=request.prompt)
            response_data = process_response(response, "both")
            response_data["message_type"] = "both"
        response_data["session_id"] = session_id
        print(f"✅ Generated response for session {session_id}")
        return response_data
    except Exception as e:
        print(f"❌ Error in generate_response: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Generation error: {str(e)}")

@router.post("/clear_history/")
async def clear_history(request: HistoryRequest):
    if request.session_id in conversation_memories:
        conversation_memories[request.session_id].clear()
        return {"status": "success", "message": "Conversation history cleared"}
    return {"status": "not_found", "message": "Session ID not found"}

@router.post("/get_history/")
async def get_history(request: HistoryRequest):
    if request.session_id in conversation_memories:
        memory = conversation_memories[request.session_id]
        messages = memory.chat_memory.messages
        history = []
        for msg in messages:
            if isinstance(msg, HumanMessage):
                history.append(f"Human: {msg.content}")
            elif isinstance(msg, AIMessage):
                history.append(f"AI: {msg.content}")
        return {"status": "success", "history": history}
    return {"status": "not_found", "message": "Session ID not found"}

@router.get("/sessions/")
async def get_active_sessions():
    return {
        "active_sessions": list(conversation_memories.keys()),
        "total_sessions": len(conversation_memories)
    }

@router.delete("/sessions/{session_id}")
async def delete_session(session_id: str):
    if session_id in conversation_memories:
        del conversation_memories[session_id]
        return {"status": "success", "message": f"Session {session_id} deleted"}
    return {"status": "not_found", "message": "Session ID not found"}

@router.get("/health")
async def health_check():
    return {
        "status": "healthy",
        "model": llm.model if llm else None,
        "active_sessions": len(conversation_memories),
        "langchain_version": "0.0.350"
    }