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

app = FastAPI()

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

# Global State
agent = None

class ChatRequest(BaseModel):
    message: str
    mode: str = "text" # text, code, etc.

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

@app.post("/api/chat")
async def chat(request: ChatRequest):
    global agent
    agent_instance = get_agent()
    
    if not agent_instance:
         # Try to load again appropriately
         agent_instance = get_agent()
         if not agent_instance:
             # Fallback mock if no model trained yet
             return {"response": "Model not trained yet! Please go to Settings and Train on a folder first."}
    
    # Simple prompt engineering based on mode
    prompt = request.message
    if request.mode == "code":
        prompt = "CODE:\n" + prompt
    
    response = agent_instance.generate(prompt, max_new_tokens=200, temperature=0.7)
    return {"response": response}

class GrammarCheckRequest(BaseModel):
    text: str

@app.post("/api/grammar/check")
async def check_grammar(request: GrammarCheckRequest):
    """
    Check grammar using SpaCy-based Grammar Checker.
    Returns POS tags, dependencies, entities, and grammar issues.
    """
    try:
        from agent.grammar_checker import get_grammar_checker
        checker = get_grammar_checker()
        result = checker.analyze(request.text)
        return {"status": "success", "analysis": result}
    except Exception as e:
        return {"status": "error", "message": str(e)}

@app.post("/api/grammar/quick")
async def quick_grammar_check(request: GrammarCheckRequest):
    """Quick grammar check returning a formatted string."""
    try:
        from agent.grammar_checker import get_grammar_checker
        checker = get_grammar_checker()
        result = checker.check_sentence(request.text)
        return {"status": "success", "result": result}
    except Exception as e:
        return {"status": "error", "message": str(e)}

# Mount static files (Frontend)
app.mount("/", StaticFiles(directory="app/static", html=True), name="static")