from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional, Dict
import os
import json
from datetime import datetime
from langchain_core.messages import HumanMessage, AIMessage
from socratic_graph import learner_app, get_resources, get_text_content
from ingest import process_new_files
import shutil
from fastapi import UploadFile, File

from fastapi.staticfiles import StaticFiles

app = FastAPI()

# Enable CORS for React development
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], # In production, restrict this to your frontend URL
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# --- MODELS ---

class Message(BaseModel):
    role: str # "user" or "assistant"
    content: str

class ChatRequest(BaseModel):
    messages: List[Message]
    grade: str
    subject: str
    topic: str
    hint_level: int
    status: str

class ChatResponse(BaseModel):
    message: Message
    hint_level: int
    status: str
    context: str

# --- UTILS ---

def convert_to_langchain(messages: List[Message]):
    lc_messages = []
    for m in messages:
        if m.role == "user":
            lc_messages.append(HumanMessage(content=m.content))
        else:
            lc_messages.append(AIMessage(content=m.content))
    return lc_messages

# --- ENDPOINTS ---

@app.get("/api/curriculum")
async def get_curriculum():
    try:
        vector_store = get_resources()
        all_data = vector_store.get(include=['metadatas'])
        metadatas = all_data['metadatas']
        
        options = {}
        for meta in metadatas:
            grade = meta.get('grade', 'Unknown')
            subject = meta.get('subject', 'Unknown')
            topics_str = meta.get('topics', 'General')
            
            if grade not in options: options[grade] = {}
            if subject not in options[grade]: options[grade][subject] = set()
                
            t_list = [t.strip().title() for t in topics_str.split(',') if t.strip()]
            for t in t_list:
                options[grade][subject].add(t)
                
        # Convert sets to sorted lists
        final_options = {}
        for g in options:
            final_options[g] = {}
            for s in options[g]:
                final_options[g][s] = sorted(list(options[g][s]))
        
        return final_options
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    try:
        lc_messages = convert_to_langchain(request.messages)
        
        initial_state = {
            "messages": lc_messages,
            "hint_level": request.hint_level,
            "context": "",
            "status": request.status,
            "safety_status": "PASS",
            "current_topic": request.topic,
            "grade": request.grade,
            "subject": request.subject,
            "selected_topic": request.topic
        }
        
        final_state = await learner_app.ainvoke(initial_state)
        response_msg = final_state["messages"][-1]
        response_text = get_text_content(response_msg.content)
        
        # Log for evaluation (consistent with app.py logic)
        log_entry = {
            "question": request.messages[-1].content if request.messages else "",
            "answer": response_text,
            "contexts": [final_state.get("context", "")],
            "timestamp": datetime.now().isoformat()
        }
        with open("eval_logs.jsonl", "a", encoding="utf-8") as f:
            f.write(json.dumps(log_entry) + "\n")
            
        return ChatResponse(
            message=Message(role="assistant", content=response_text),
            hint_level=final_state["hint_level"],
            status=final_state["status"],
            context=final_state.get("context", "")
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/ingest")
async def ingest():
    try:
        result = process_new_files()
        return {"status": "success", "message": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/upload")
async def upload_file(grade: str, subject: str, file: UploadFile = File(...)):
    try:
        # Create directory structure: data/Grade/Subject/
        target_dir = os.path.join("data", grade, subject)
        os.makedirs(target_dir, exist_ok=True)
        
        file_path = os.path.join(target_dir, file.filename)
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
            
        return {"status": "success", "filename": file.filename}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Mount the React production build output directory if it exists
frontend_dist_dir = os.path.join(os.path.dirname(__file__), "frontend", "dist")
if os.path.exists(frontend_dist_dir):
    app.mount("/", StaticFiles(directory=frontend_dist_dir, html=True), name="static")
else:
    # Optional fallback or message
    print("Warning: frontend/dist directory not found. Static files will not be served.")

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