from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import onnxruntime as ort
import numpy as np
import chess
import time
import logging
import os
from typing import Optional, Tuple

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class NexusNanoEngine:
    PIECE_VALUES = {chess.PAWN: 1, chess.KNIGHT: 3, chess.BISHOP: 3, chess.ROOK: 5, chess.QUEEN: 9, chess.KING: 0}
    
    def __init__(self, model_path: str):
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model not found: {model_path}")
        logger.info(f"Loading model from {model_path}...")
        sess_options = ort.SessionOptions()
        sess_options.intra_op_num_threads = 2
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.session = ort.InferenceSession(model_path, sess_options=sess_options, providers=['CPUExecutionProvider'])
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        self.nodes = 0
        logger.info("✅ Nexus-Nano engine loaded")
    
    def fen_to_tensor(self, fen: str) -> np.ndarray:
        board = chess.Board(fen)
        tensor = np.zeros((1, 12, 8, 8), dtype=np.float32)
        piece_map = {chess.PAWN: 0, chess.KNIGHT: 1, chess.BISHOP: 2, chess.ROOK: 3, chess.QUEEN: 4, chess.KING: 5}
        for sq, piece in board.piece_map().items():
            r, f = divmod(sq, 8)
            ch = piece_map[piece.piece_type] + (6 if piece.color == chess.BLACK else 0)
            tensor[0, ch, r, f] = 1.0
        return tensor
    
    def evaluate(self, board: chess.Board) -> float:
        self.nodes += 1
        tensor = self.fen_to_tensor(board.fen())
        output = self.session.run([self.output_name], {self.input_name: tensor})
        score = float(output[0][0][0]) * 400.0
        return -score if board.turn == chess.BLACK else score
    
    def order_moves(self, board, moves):
        scored = []
        for m in moves:
            s = 0
            if board.is_capture(m):
                v, a = board.piece_at(m.to_square), board.piece_at(m.from_square)
                if v and a: s = self.PIECE_VALUES.get(v.piece_type, 0) * 10 - self.PIECE_VALUES.get(a.piece_type, 0)
            if m.promotion == chess.QUEEN: s += 90
            scored.append((s, m))
        scored.sort(key=lambda x: x[0], reverse=True)
        return [m for _, m in scored]
    
    def alpha_beta(self, board, depth, alpha, beta):
        if board.is_game_over(): return (-10000 if board.is_checkmate() else 0), None
        if depth == 0: return self.evaluate(board), None
        moves = list(board.legal_moves)
        if not moves: return 0, None
        moves = self.order_moves(board, moves)
        best_move, best_score = moves[0], float('-inf')
        for move in moves:
            board.push(move)
            score, _ = self.alpha_beta(board, depth - 1, -beta, -alpha)
            score = -score
            board.pop()
            if score > best_score: best_score, best_move = score, move
            alpha = max(alpha, score)
            if alpha >= beta: break
        return best_score, best_move
    
    def search(self, fen: str, depth: int = 3):
        board = chess.Board(fen)
        self.nodes = 0
        moves = list(board.legal_moves)
        if not moves: return {'best_move': '0000', 'evaluation': 0.0, 'nodes': 0, 'depth': 0}
        if len(moves) == 1: return {'best_move': moves[0].uci(), 'evaluation': round(self.evaluate(board)/100, 2), 'nodes': 1, 'depth': 0}
        best_move, best_score, current_depth = moves[0], float('-inf'), 1
        for d in range(1, depth + 1):
            try:
                score, move = self.alpha_beta(board, d, float('-inf'), float('inf'))
                if move: best_move, best_score, current_depth = move, score, d
            except: break
        return {'best_move': best_move.uci(), 'evaluation': round(best_score/100, 2), 'depth': current_depth, 'nodes': self.nodes}

app = FastAPI(title="Nexus-Nano API", version="1.0.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])

engine = None

class MoveRequest(BaseModel):
    fen: str
    depth: Optional[int] = Field(3, ge=1, le=5)

class MoveResponse(BaseModel):
    best_move: str
    evaluation: float
    depth_searched: int
    nodes_evaluated: int
    time_taken: int

@app.on_event("startup")
async def startup():
    global engine
    logger.info("🚀 Starting Nexus-Nano API...")
    model_path = "/app/nexus_nano.onnx"
    try:
        engine = NexusNanoEngine(model_path)
        logger.info("✅ Engine ready")
    except Exception as e:
        logger.error(f"❌ Failed to load engine: {e}")
        raise

@app.get("/health")
async def health():
    return {"status": "healthy" if engine else "unhealthy", "model_loaded": engine is not None, "version": "1.0.0"}

@app.post("/get-move", response_model=MoveResponse)
async def get_move(req: MoveRequest):
    if not engine: raise HTTPException(503, "Engine not loaded")
    try: chess.Board(req.fen)
    except: raise HTTPException(400, "Invalid FEN")
    start = time.time()
    try:
        result = engine.search(req.fen, req.depth)
        elapsed = int((time.time() - start) * 1000)
        logger.info(f"Move: {result['best_move']} | Eval: {result['evaluation']:+.2f} | Time: {elapsed}ms")
        return MoveResponse(best_move=result['best_move'], evaluation=result['evaluation'], 
                          depth_searched=result['depth'], nodes_evaluated=result['nodes'], time_taken=elapsed)
    except Exception as e:
        logger.error(f"Error: {e}")
        raise HTTPException(500, str(e))

@app.get("/")
async def root():
    return {"name": "Nexus-Nano", "version": "1.0.0", "status": "online" if engine else "starting"}

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