import gradio as gr
import numpy as np
import cv2
from PIL import Image
import os

import chess
import chess.engine

# Robust imports
ENGINE_CANDIDATES = ["stockfish", "/usr/bin/stockfish", "/usr/local/bin/stockfish"]

try:
    from ultralytics import YOLO
    YOLO_AVAILABLE = True
except Exception as e:
    YOLO_AVAILABLE = False
    YOLO_IMPORT_ERROR = str(e)

YOLO_MODEL_ID = "yamero999/chess-piece-detection-yolo11n"
YOLO_IMGSZ = 640
YOLO_CONF = 0.35
STOCKFISH_SKILL_LEVEL = 8

LABEL_TO_FEN = {
    "wpawn": "P", "wknight": "N", "wbishop": "B", "wrook": "R", "wqueen": "Q", "wking": "K",
    "bpawn": "p", "bknight": "n", "bbishop": "b", "brook": "r", "bqueen": "q", "bking": "k",
    "white_pawn": "P","white_knight":"N","white_bishop":"B","white_rook":"R","white_queen":"Q","white_king":"K",
    "black_pawn":"p","black_knight":"n","black_bishop":"b","black_rook":"r","black_queen":"q","black_king":"k",
}

def load_yolo():
    if not YOLO_AVAILABLE:
        raise RuntimeError(f"Ultralytics not available: {YOLO_IMPORT_ERROR}")
    try:
        model = YOLO(YOLO_MODEL_ID)
        return model
    except Exception as e:
        raise RuntimeError(f"Failed to load YOLO model '{YOLO_MODEL_ID}': {e}")

def detect_board_corners(img_bgr):
    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (5,5), 0)
    thr = cv2.adaptiveThreshold(gray,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY_INV, 31, 5)
    kernel = np.ones((3,3), np.uint8)
    thr = cv2.morphologyEx(thr, cv2.MORPH_CLOSE, kernel, iterations=2)

    contours, _ = cv2.findContours(thr, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None
    contours = sorted(contours, key=cv2.contourArea, reverse=True)

    for cnt in contours[:5]:
        peri = cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)
        if len(approx) == 4:
            pts = approx.reshape(4,2).astype(np.float32)
            s = pts.sum(axis=1)
            diff = np.diff(pts, axis=1).reshape(-1)
            tl = pts[np.argmin(s)]
            br = pts[np.argmax(s)]
            tr = pts[np.argmin(diff)]
            bl = pts[np.argmax(diff)]
            return np.array([tl,tr,br,bl], dtype=np.float32)
    return None

def warp_to_topdown(img_bgr, corners):
    dst = np.array([[0,0],[800,0],[800,800],[0,800]], dtype=np.float32)
    M = cv2.getPerspectiveTransform(corners, dst)
    warped = cv2.warpPerspective(img_bgr, M, (800,800))
    return warped

def yolo_detect_pieces(model, img_bgr):
    results = model.predict(source=img_bgr[...,::-1], imgsz=YOLO_IMGSZ, conf=YOLO_CONF, verbose=False)
    dets = []
    if not results:
        return dets
    res = results[0]
    names = res.names
    for b in res.boxes:
        cls_id = int(b.cls[0].item())
        conf = float(b.conf[0].item())
        x1,y1,x2,y2 = b.xyxy[0].tolist()
        dets.append({
            "label": names.get(cls_id, str(cls_id)).lower(),
            "conf": conf,
            "bbox": (float(x1), float(y1), float(x2), float(y2))
        })
    return dets

def piece_square_mapping(warped_bgr, detections):
    mapping = {}
    square_size = 100
    for det in detections:
        fen_letter = LABEL_TO_FEN.get(det["label"])
        if not fen_letter:
            # coarse fallback based on substrings
            lbl = det["label"]
            if "pawn" in lbl: fen_letter = 'P' if 'w' in lbl else 'p'
            elif "knight" in lbl: fen_letter = 'N' if 'w' in lbl else 'n'
            elif "bishop" in lbl: fen_letter = 'B' if 'w' in lbl else 'b'
            elif "rook" in lbl: fen_letter = 'R' if 'w' in lbl else 'r'
            elif "queen" in lbl: fen_letter = 'Q' if 'w' in lbl else 'q'
            elif "king" in lbl: fen_letter = 'K' if 'w' in lbl else 'k'
        if not fen_letter:
            continue
        x1,y1,x2,y2 = det["bbox"]
        cx = (x1 + x2) / 2.0
        cy = (y1 + y2) / 2.0
        col = int(np.clip(cx // square_size, 0, 7))
        row = int(np.clip(cy // square_size, 0, 7))
        file_char = chr(ord('a') + col)
        rank_char = str(8 - row)
        sq = f"{file_char}{rank_char}"
        old = mapping.get(sq)
        if old is None or det["conf"] > old["conf"]:
            mapping[sq] = {"fen": fen_letter, "conf": det["conf"]}
    return {k: v["fen"] for k,v in mapping.items()}

def mapping_to_fen(square_map):
    rows = []
    for r in range(8, 0, -1):
        row_str, empty = "", 0
        for c in range(8):
            sq = f"{chr(ord('a')+c)}{r}"
            if sq in square_map:
                if empty: row_str += str(empty); empty = 0
                row_str += square_map[sq]
            else:
                empty += 1
        if empty: row_str += str(empty)
        rows.append(row_str)
    board_fen = "/".join(rows)
    return f"{board_fen} w - - 0 1"

def find_stockfish():
    for path in ENGINE_CANDIDATES:
        try:
            eng = chess.engine.SimpleEngine.popen_uci(path)
            return eng
        except Exception:
            continue
    raise RuntimeError("Stockfish engine not found. Ensure apt.txt installs it or set ENGINE_CANDIDATES.")

def san_best_move_and_reason(fen):
    try:
        engine = find_stockfish()
    except Exception as e:
        return None, f"Engine error: {e}"
    try:
        try:
            engine.configure({"Skill Level": int(STOCKFISH_SKILL_LEVEL)})
        except Exception:
            pass
        board = chess.Board(fen)
        info0 = engine.analyse(board, chess.engine.Limit(depth=10))
        result = engine.play(board, chess.engine.Limit(depth=12))
        move = result.move
        board.push(move)
        info1 = engine.analyse(board, chess.engine.Limit(depth=10))

        def cp(info):
            s = info.get("score")
            if s is None: return None
            try:
                return s.white().score(mate_score=100000)
            except Exception:
                return None

        before, after = cp(info0), cp(info1)
        san = board.peek().san()
        board.pop()

        text = "This move improves your coordination and keeps the position stable."
        if before is not None and after is not None:
            delta = after - before
            if delta >= 80:
                text = "A strong move that clearly improves your position and creates threats."
            elif delta >= 30:
                text = "A good developing move that gains a small but steady advantage."
            elif delta >= 5:
                text = "A useful move that slightly improves your position."
            elif delta >= -5:
                text = "A safe, solid move that keeps the balance."
            else:
                text = "A practical choice to avoid complications."
        return san, text
    finally:
        try:
            engine.quit()
        except Exception:
            pass

def try_opening_name(fen):
    try:
        import chess.openings
        board = chess.Board(fen)
        name = chess.openings.opening_name(board)
        return name or "No named opening (or midgame)."
    except Exception:
        return "No named opening (or midgame)."

def draw_move(img_bgr, san, fen):
    try:
        board = chess.Board(fen)
        mv = board.parse_san(san)
    except Exception:
        return img_bgr
    def center(sq):
        file = chess.square_file(sq)
        rank = chess.square_rank(sq)
        col = file
        row_top = 7 - rank
        return (int((col+0.5)*100), int((row_top+0.5)*100))
    a = center(mv.from_square); b = center(mv.to_square)
    out = img_bgr.copy()
    cv2.arrowedLine(out, a, b, (0,255,0), 4, tipLength=0.25)
    return out

def process(image):
    if image is None:
        return None, "", "", "", "Please upload a board image."
    img_bgr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
    corners = detect_board_corners(img_bgr)
    warped = cv2.resize(img_bgr, (800,800)) if corners is None else warp_to_topdown(img_bgr, corners)

    # Load YOLO
    try:
        model = load_yolo()
    except Exception as e:
        return Image.fromarray(cv2.cvtColor(warped, cv2.COLOR_BGR2RGB)), "", "", "", f"Detector load error: {e}"

    try:
        dets = yolo_detect_pieces(model, warped)
        sq_map = piece_square_mapping(warped, dets)
        fen = mapping_to_fen(sq_map)
    except Exception as e:
        return Image.fromarray(cv2.cvtColor(warped, cv2.COLOR_BGR2RGB)), "", "", "", f"Detection/FEN error: {e}"

    san, why = san_best_move_and_reason(fen)
    if san is None:
        out = warped
        move = ""
        err = why
    else:
        out = draw_move(warped, san, fen)
        move = san
        err = ""

    opening = try_opening_name(fen)
    out_img = Image.fromarray(cv2.cvtColor(out, cv2.COLOR_BGR2RGB))
    return out_img, fen, move, opening, why if not err else err

title_md = "# Chess Assist (v3): Image → FEN → Best Move"

with gr.Blocks() as demo:
    gr.Markdown(title_md)
    with gr.Row():
        with gr.Column():
            img = gr.Image(type="pil", label="Upload a chessboard photo")
            go = gr.Button("Analyze")
        with gr.Column():
            vis = gr.Image(type="pil", label="Move overlay")
            fen = gr.Textbox(label="FEN")
            mv = gr.Textbox(label="Best Move (SAN)")
            opening = gr.Textbox(label="Opening")
            why = gr.Textbox(label="Why this move? / Errors")
    go.click(process, inputs=[img], outputs=[vis, fen, mv, opening, why])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)