import os
import chess
import chess.engine
import yaml
import platform
import asyncio
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
from typing import List
from contextlib import asynccontextmanager

# Global engine instances
rust_engine_pool = None
veloct_engine_pool = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global rust_engine_pool, veloct_engine_pool
    
    # Initialize the Rust Engine globally on startup!
    engine_name = "Neurex.exe" if platform.system() == "Windows" else "Neurex"
    engine_path = os.path.join(os.path.dirname(__file__), "Neurex_Engine", engine_name)
    
    if os.path.exists(engine_path):
        print(f"Booting up {engine_name} globally...")
        try:
            _, rust_engine_pool = await asyncio.wait_for(
                chess.engine.popen_uci(engine_path), timeout=60.0
            )
            print("Neurex Engine initialized successfully!")
        except Exception as e:
            print(f"Failed to boot Neurex engine: {e}")
            
    veloct_name = "VeloCT_Ultimate.exe" if platform.system() == "Windows" else "VeloCT_Ultimate"
    veloct_path = os.path.join(os.path.dirname(__file__), veloct_name)
    if os.path.exists(veloct_path):
        print(f"Booting up {veloct_name} globally...")
        try:
            _, veloct_engine_pool = await asyncio.wait_for(
                chess.engine.popen_uci(veloct_path), timeout=60.0
            )
            print("VeloCT Engine initialized successfully!")
        except Exception as e:
            print(f"Failed to boot VeloCT engine: {e}")
            
    yield # App is running
    
    # Shutdown
    if rust_engine_pool:
        await rust_engine_pool.quit()
    if veloct_engine_pool:
        await veloct_engine_pool.quit()

app = FastAPI(lifespan=lifespan)
templates = Jinja2Templates(directory="templates")

class MoveRequest(BaseModel):
    history: List[str] = ["<bos>"]
    fen: str = chess.STARTING_FEN
    engine_choice: str = "neural"
    time_limit: float = 5.0

class PGNRequest(BaseModel):
    pgn: str

@app.post("/save_pgn")
async def save_pgn(req: PGNRequest):
    import datetime
    os.makedirs("games", exist_ok=True)
    filename = f"games/game_{int(datetime.datetime.now().timestamp())}.pgn"
    with open(filename, "w") as f:
        f.write(req.pgn)
    return {"status": "ok", "file": filename}

@app.get("/")
async def index(request: Request):
    return templates.TemplateResponse(request=request, name="index.html")

@app.post("/move")
async def make_move(req: MoveRequest):
    history = req.history
    engine_choice = req.engine_choice
    
    # Reconstruct board from history
    board = chess.Board()
    for m in history:
        if m != "<bos>":
            try:
                board.push_uci(m)
            except:
                pass
    if board.is_game_over():
        return JSONResponse({'error': 'Game Over'})
        
    ai_eval = 0.0
    ai_move = None
    
    try:
        if engine_choice == 'veloct':
            if veloct_engine_pool is None:
                print("ERROR: VeloCT engine is not loaded.")
                return JSONResponse({'ai_move': None, 'eval': 0.0})
                
            result = await veloct_engine_pool.play(board, chess.engine.Limit(time=0.5))
            ai_move = result.move.uci()
            
            info = await veloct_engine_pool.analyse(board, chess.engine.Limit(time=0.1))
            if "score" in info:
                score = info["score"].pov(chess.WHITE)
                if score.is_mate():
                    ai_eval = 1.0 if score.mate() > 0 else -1.0
                else:
                    ai_eval = max(-1.0, min(1.0, score.score() / 500.0))
        else:
            # Use the globally initialized Rust Engine
            if rust_engine_pool is None:
                print("ERROR: Rust engine is not loaded.")
                return JSONResponse({'ai_move': None, 'eval': 0.0})
                
            result = await rust_engine_pool.play(board, chess.engine.Limit(time=req.time_limit), info=chess.engine.INFO_ALL)
            ai_move = result.move.uci()
            
            if "score" in result.info:
                score = result.info["score"].white()
                if score.is_mate():
                    ai_eval = 1.0 if score.mate() > 0 else -1.0
                else:
                    ai_eval = max(-1.0, min(1.0, score.score() / 1000.0))
            
            engine_info = {}
            if "depth" in result.info: engine_info["depth"] = result.info["depth"]
            if "nodes" in result.info: engine_info["nodes"] = result.info["nodes"]
            if "nps" in result.info: engine_info["nps"] = result.info["nps"]
                    
    except Exception as e:
        import traceback
        traceback.print_exc()
        print(f"Engine error: {repr(e)}")
        ai_move = None
        engine_info = {}
            
    return JSONResponse({'ai_move': ai_move, 'eval': ai_eval, 'engine_info': engine_info})

if __name__ == '__main__':
    import uvicorn
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run("app:app", host="0.0.0.0", port=port, log_level="info")

# Force push