"""
Nexus-Nano Inference API
Ultra-lightweight single-file engine
No modular architecture - pure speed optimization
"""

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
from pathlib import Path
from typing import Optional, Tuple

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ==================== NANO ENGINE (Single File) ====================

class NexusNanoEngine:
    """
    Ultra-lightweight chess engine
    Pure alpha-beta, no cache, minimal overhead
    """
    
    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):
        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 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: chess.Board, moves):
        """Simple MVV-LVA ordering"""
        scored = []
        for m in moves:
            s = 0
            if board.is_capture(m):
                v = board.piece_at(m.to_square)
                a = board.piece_at(m.from_square)
                if v and a:
                    s = self.PIECE_VALUES.get(v.piece_type, 0) * 10
                    s -= 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: chess.Board, 
        depth: int, 
        alpha: float, 
        beta: float
    ) -> Tuple[float, Optional[chess.Move]]:
        
        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 = moves[0]
        best_score = 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 = score
                best_move = 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 len(moves) == 0:
            return {'best_move': '0000', 'evaluation': 0.0, 'nodes': 0}
        if len(moves) == 1:
            return {
                'best_move': moves[0].uci(),
                'evaluation': round(self.evaluate(board) / 100.0, 2),
                'nodes': 1
            }
        
        best_move = moves[0]
        best_score = float('-inf')
        
        for d in range(1, depth + 1):
            try:
                score, move = self.alpha_beta(board, d, float('-inf'), float('inf'))
                if move:
                    best_move = move
                    best_score = score
            except:
                break
        
        return {
            'best_move': best_move.uci(),
            'evaluation': round(best_score / 100.0, 2),
            'depth': d,
            'nodes': self.nodes
        }


# ==================== FASTAPI APP ====================

app = FastAPI(
    title="Nexus-Nano API",
    description="Ultra-lightweight chess engine",
    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...")
    try:
        engine = NexusNanoEngine("/app/models/nexus_nano.onnx")
    except Exception as e:
        logger.error(f"❌ Failed: {e}")
        raise


@app.get("/health")
async def health():
    return {"status": "healthy", "model": "nexus-nano", "version": "1.0.0"}


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


@app.get("/")
async def root():
    return {
        "name": "Nexus-Nano API",
        "version": "1.0.0",
        "model": "2.8M parameters",
        "speed": "Lightning-fast"
    }


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