from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
import os
import json
import time
from datetime import datetime
from agent.inference import InferenceEngine
from model.config import ModelConfig
import torch

app = FastAPI()

# Enable CORS (Critical for cross-port communication)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global State
agent = None
CONVERSATIONS_FILE = os.path.join("data", "conversations.json")

# Ensure data directory exists
os.makedirs("data", exist_ok=True)
if not os.path.exists(CONVERSATIONS_FILE):
    with open(CONVERSATIONS_FILE, "w") as f:
        json.dump({"sessions": {}}, f)

class ChatRequest(BaseModel):
    message: str
    session_id: str = "default"
    mode: str = "text"

def get_agent():
    global agent
    if agent is None or agent.model is None:
        model_path = "sail.pt"
        if os.path.exists(model_path):
            try:
                agent = InferenceEngine(model_path=model_path)
                print(f"Agent Model Loaded from {model_path}.")
            except Exception as e:
                print(f"Failed to load agent: {e}")
    return agent

def log_conversation(session_id, role, content):
    try:
        with open(CONVERSATIONS_FILE, "r") as f:
            data = json.load(f)
        
        if session_id not in data["sessions"]:
            data["sessions"][session_id] = {
                "id": session_id,
                "timestamp": datetime.now().isoformat(),
                "messages": []
            }
        
        data["sessions"][session_id]["messages"].append({
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        })
        
        # Update preview (first user message)
        if role == "user" and len(data["sessions"][session_id]["messages"]) <= 2:
            data["sessions"][session_id]["preview"] = content[:50] + ("..." if len(content) > 50 else "")

        with open(CONVERSATIONS_FILE, "w") as f:
            json.dump(data, f, indent=2)
    except Exception as e:
        print(f"Failed to log conversation: {e}")

@app.get("/api/chat/status")
async def get_status():
    return {"status": "connected", "model": "sail.pt" if os.path.exists("sail.pt") else "none"}

@app.post("/api/chat")
async def chat(request: ChatRequest):
    agent_instance = get_agent()
    
    if not agent_instance:
         return {"response": "Model not loaded! Please ensure sail.pt exists in the root directory."}
    
    # Log user message
    log_conversation(request.session_id, "user", request.message)
    
    # Generate response
    response = agent_instance.generate(request.message, max_new_tokens=256, temperature=0.7)
    
    # Log system response
    log_conversation(request.session_id, "assistant", response)
    
    return {"response": response}

@app.get("/api/history")
async def get_history():
    try:
        with open(CONVERSATIONS_FILE, "r") as f:
            data = json.load(f)
        sessions = []
        for s_id, s_data in data["sessions"].items():
            sessions.append({
                "id": s_id,
                "timestamp": s_data["timestamp"],
                "preview": s_data.get("preview", "Untitled Chat")
            })
        # Sort by latest
        sessions.sort(key=lambda x: x["timestamp"], reverse=True)
        return {"sessions": sessions}
    except Exception as e:
        return {"sessions": [], "error": str(e)}

@app.get("/api/history/{session_id}")
async def get_session_history(session_id: str):
    try:
        with open(CONVERSATIONS_FILE, "r") as f:
            data = json.load(f)
        if session_id in data["sessions"]:
            return {"messages": data["sessions"][session_id]["messages"]}
        return {"messages": []}
    except Exception as e:
        return {"messages": [], "error": str(e)}

@app.post("/api/grammar/check")
async def check_grammar(request: dict):
    # Simplified placeholder to keep API compatible if needed
    return {"status": "success", "analysis": "Grammar check currently disabled in simplified mode."}