#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Lightweight image-to-text API for Jetson Nano + EdgeTPU.
Compatible with Python 3.6.
"""
import argparse
import base64
import io
import os
import sys
import time
import uuid
import threading
from flask import Flask, request, jsonify

from core import image_io, coral_vision, text_detect, ocr, schema

try:
    import cv2
except ImportError:  # pragma: no cover
    cv2 = None
import numpy as np
from PIL import Image, ImageDraw

# ---- Configuration ----

def env_bool(name, default):
    val = os.getenv(name)
    if val is None:
        return default
    return val.lower() in ("1", "true", "yes", "on")

def env_float(name, default):
    val = os.getenv(name)
    try:
        return float(val) if val is not None else default
    except ValueError:
        return default

def env_int(name, default):
    val = os.getenv(name)
    try:
        return int(val) if val is not None else default
    except ValueError:
        return default

MODEL_EDGETPU_PATH = os.getenv("MODEL_EDGETPU_PATH", "./models/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite")
MODEL_LABELS_PATH = os.getenv("MODEL_LABELS_PATH", "./models/coco_labels.txt")
MODEL_EAST_PATH = os.getenv("MODEL_EAST_PATH", "./models/frozen_east_text_detection.pb")
OCR_LANG = os.getenv("OCR_LANG", "eng+rus")
MAX_IMAGE_MB = env_int("MAX_IMAGE_MB", 10)
CONF_THRESHOLD = env_float("CONF_THRESHOLD", 0.4)
RETURN_BOXES = env_bool("RETURN_BOXES", True)
OCR_CONFIG = os.getenv("OCR_CONFIG", "--oem 1 --psm 6")
TESSDATA_PREFIX = os.getenv("TESSDATA_PREFIX")
USE_CUDA_DNN = env_bool("USE_CUDA_DNN", False)
DEBUG_VISUAL = env_bool("DEBUG_VISUAL", False)
SHOW_GUI = env_bool("SHOW_GUI", False)
GUI_TARGET_WIDTH = env_int("GUI_TARGET_WIDTH", 1600)
OCR_UPSCALE = env_float("OCR_UPSCALE", 1.5)  # 1.0 = no upscale
OCR_CLAHE = env_bool("OCR_CLAHE", True)
OCR_BILATERAL = env_bool("OCR_BILATERAL", False)
OCR_BLUR = env_int("OCR_BLUR", 0)  # 0 disables gaussian blur
OCR_INVERT = env_bool("OCR_INVERT", False)
OCR_USE_THRESHOLD = env_bool("OCR_USE_THRESHOLD", True)
OCR_AUTO_LANG = env_bool("OCR_AUTO_LANG", True)
MAX_CONTENT_LENGTH = MAX_IMAGE_MB * 1024 * 1024

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = MAX_CONTENT_LENGTH

if not TESSDATA_PREFIX:
    # Attempt common system locations
    for candidate in (
        "/usr/share/tesseract-ocr/4.00/tessdata_best",
        "/usr/share/tesseract-ocr/4.00/tessdata",
    ):
        if os.path.exists(candidate):
            TESSDATA_PREFIX = candidate
            break

if TESSDATA_PREFIX:
    os.environ["TESSDATA_PREFIX"] = TESSDATA_PREFIX

_GUI_FRAME = None
_GUI_LOCK = threading.Lock()
_GUI_THREAD = None

def _start_gui_thread():
    if not SHOW_GUI or cv2 is None:
        return
    global _GUI_THREAD
    if _GUI_THREAD is not None:
        return
    def _loop():
        cv2.namedWindow("image2text debug", cv2.WINDOW_NORMAL)
        while True:
            frame = None
            with _GUI_LOCK:
                if _GUI_FRAME is not None:
                    frame = _GUI_FRAME.copy()
            if frame is not None:
                cv2.imshow("image2text debug", frame)
                cv2.resizeWindow("image2text debug", frame.shape[1], frame.shape[0])
            cv2.waitKey(50)
    _GUI_THREAD = threading.Thread(target=_loop, daemon=True)
    _GUI_THREAD.start()

def _show_placeholder():
    if not SHOW_GUI or cv2 is None:
        return
    _gui_show_message("image2text waiting for request")


def _gui_show(frame):
    if not SHOW_GUI or cv2 is None or frame is None:
        return
    global _GUI_FRAME
    with _GUI_LOCK:
        _GUI_FRAME = frame


def _gui_show_message(msg):
    if not SHOW_GUI or cv2 is None:
        return
    width = GUI_TARGET_WIDTH if GUI_TARGET_WIDTH > 0 else 1280
    height = int(width * 9 / 16)
    canvas = np.ones((height, width, 3), dtype=np.uint8) * 240
    cv2.putText(canvas, msg[:60], (40, height // 2),
                cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 2, cv2.LINE_AA)
    _gui_show(canvas)

# ---- Load models once ----

coral = coral_vision.CoralVision(
    model_path=MODEL_EDGETPU_PATH,
    labels_path=MODEL_LABELS_PATH,
    score_threshold=CONF_THRESHOLD,
)

txt_detector = text_detect.TextDetector(
    east_model_path=MODEL_EAST_PATH,
    score_threshold=CONF_THRESHOLD,
    use_cuda=USE_CUDA_DNN,
)

# ---- Helpers ----

def parse_image_from_request(req):
    if 'image' in req.files:
        data = req.files['image'].read()
        source = 'multipart'
    else:
        try:
            payload = req.get_json(force=True, silent=False)
        except Exception:
            payload = None
        if payload and 'image_base64' in payload:
            b64_str = payload['image_base64'] or ''
            if ',' in b64_str:
                b64_str = b64_str.split(',', 1)[1]
            try:
                data = base64.b64decode(b64_str, validate=False)
            except Exception:
                return None, 'invalid_base64'
            source = 'base64'
        else:
            return None, 'not_found'
    if not data:
        return None, 'empty'
    if len(data) > MAX_CONTENT_LENGTH:
        return None, 'too_large'
    return data, source


def build_debug_image(img_rgb, boxes, blocks, objects, target_width=1024, return_image=False):
    """Return (base64 JPEG, optional ndarray for display) with drawn boxes."""
    try:
        vis = img_rgb.copy()
        h, w = vis.shape[:2]
        scale = 1.0
        if target_width and target_width > 0 and w != target_width:
            scale = float(target_width) / float(w)
            # avoid crazy upscales
            if scale > 3.0:
                scale = 3.0
            new_w = int(w * scale)
            new_h = int(h * scale)
            vis = np.array(Image.fromarray(vis).resize((new_w, new_h)))
        # draw using cv2 if available
        if cv2 is not None:
            vis_bgr = cv2.cvtColor(vis, cv2.COLOR_RGB2BGR)
            # text boxes in blue
            for b in boxes:
                x1, y1, x2, y2 = [int(v * scale) for v in b]
                cv2.rectangle(vis_bgr, (x1, y1), (x2, y2), (255, 0, 0), 2)
            for blk in blocks:
                x1, y1, x2, y2 = [int(v * scale) for v in blk.get("bbox", [0, 0, 0, 0])]
                cv2.rectangle(vis_bgr, (x1, y1), (x2, y2), (0, 128, 255), 2)
                txt = blk.get("text", "")[:20]
                cv2.putText(vis_bgr, txt, (x1, max(0, y1 - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 128, 255), 1, cv2.LINE_AA)
            # objects in green
            for obj in objects or []:
                bbox = obj.get("bbox", [0, 0, 0, 0])
                x1, y1, x2, y2 = [int(v * scale) for v in bbox]
                cv2.rectangle(vis_bgr, (x1, y1), (x2, y2), (0, 200, 0), 2)
                lbl = obj.get("label", "obj")
                cv2.putText(vis_bgr, lbl, (x1, max(0, y1 - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 200, 0), 1, cv2.LINE_AA)
            ok, buf = cv2.imencode(".jpg", vis_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
            if not ok:
                return None, None
            b64 = base64.b64encode(buf.tobytes()).decode("ascii")
            return b64, vis_bgr if return_image else None
        # fallback to PIL
        img = Image.fromarray(vis)
        draw = ImageDraw.Draw(img)
        for b in boxes:
            x1, y1, x2, y2 = [int(v * scale) for v in b]
            draw.rectangle([x1, y1, x2, y2], outline="blue", width=2)
        for blk in blocks:
            x1, y1, x2, y2 = [int(v * scale) for v in blk.get("bbox", [0, 0, 0, 0])]
            draw.rectangle([x1, y1, x2, y2], outline="orange", width=2)
            txt = blk.get("text", "")[:20]
            draw.text((x1, max(0, y1 - 10)), txt, fill="orange")
        for obj in objects or []:
            bbox = obj.get("bbox", [0, 0, 0, 0])
            x1, y1, x2, y2 = [int(v * scale) for v in bbox]
            draw.rectangle([x1, y1, x2, y2], outline="green", width=2)
            draw.text((x1, max(0, y1 - 10)), obj.get("label", "obj"), fill="green")
        buf = io.BytesIO()
        img.save(buf, format="JPEG", quality=85)
        b64 = base64.b64encode(buf.getvalue()).decode("ascii")
        return b64, np.array(img) if return_image else None
    except Exception:
        return None, None


def build_text_gui_image(img_rgb, boxes, blocks, target_width=1600):
    """Build a text-centric GUI frame: regions on left, OCR lines on right."""
    if cv2 is None or img_rgb is None:
        return None
    try:
        src = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
        h, w = src.shape[:2]
        frame_w = target_width if target_width and target_width > 0 else 1600
        if frame_w < 900:
            frame_w = 900

        left_w = int(frame_w * 0.64)
        scale = float(left_w) / float(max(1, w))
        left_h = max(360, int(h * scale))
        left = cv2.resize(src, (left_w, left_h), interpolation=cv2.INTER_AREA if scale < 1.0 else cv2.INTER_CUBIC)

        # text detector regions
        for idx, b in enumerate(boxes or []):
            try:
                x1, y1, x2, y2 = [int(v * scale) for v in b]
            except Exception:
                continue
            cv2.rectangle(left, (x1, y1), (x2, y2), (255, 90, 0), 2)
            cv2.putText(left, "T{}".format(idx + 1), (x1, max(12, y1 - 4)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 90, 0), 1, cv2.LINE_AA)

        # OCR blocks (where text was actually extracted)
        for idx, blk in enumerate(blocks or []):
            try:
                x1, y1, x2, y2 = [int(v * scale) for v in blk.get("bbox", [0, 0, 0, 0])]
            except Exception:
                continue
            cv2.rectangle(left, (x1, y1), (x2, y2), (0, 160, 255), 2)
            cv2.putText(left, "R{}".format(idx + 1), (x1, min(left_h - 8, y2 + 14)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 160, 255), 1, cv2.LINE_AA)

        right_w = frame_w - left_w
        if right_w < 320:
            right_w = 320
            left = cv2.resize(left, (frame_w - right_w, left_h))
            left_w = left.shape[1]
        right = np.ones((left_h, right_w, 3), dtype=np.uint8) * 245

        y = 28
        line_h = 22
        cv2.putText(right, "OCR text panel", (14, y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (30, 30, 30), 2, cv2.LINE_AA)
        y += line_h + 4
        cv2.putText(right, "detected regions: {}".format(len(boxes or [])), (14, y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, (70, 70, 70), 1, cv2.LINE_AA)
        y += line_h
        cv2.putText(right, "recognized blocks: {}".format(len(blocks or [])), (14, y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, (70, 70, 70), 1, cv2.LINE_AA)
        y += line_h + 4
        cv2.line(right, (12, y), (right_w - 12, y), (200, 200, 200), 1)
        y += 18

        if not blocks:
            cv2.putText(right, "No OCR text blocks", (14, y),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.56, (100, 100, 100), 1, cv2.LINE_AA)
        else:
            visible = blocks[:10]
            max_chars = max(18, int((right_w - 24) / 8))
            for idx, blk in enumerate(visible):
                conf = blk.get("conf", -1)
                conf_str = "{:.1f}".format(conf) if isinstance(conf, (float, int)) and conf >= 0 else "n/a"
                header = "#{:02d} conf={}".format(idx + 1, conf_str)
                cv2.putText(right, header, (14, y), cv2.FONT_HERSHEY_SIMPLEX, 0.50, (20, 20, 20), 1, cv2.LINE_AA)
                y += line_h - 2
                raw = (blk.get("text", "") or "").replace("\n", " ").replace("\r", " ").strip()
                txt = " ".join(raw.split())
                if len(txt) > max_chars:
                    txt = txt[:max_chars - 1] + "..."
                cv2.putText(right, txt if txt else "<empty>", (14, y),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.52, (30, 30, 30), 1, cv2.LINE_AA)
                y += line_h
                if y > left_h - 18:
                    break

        frame = np.concatenate([left, right], axis=1)
        return frame
    except Exception:
        return None

# ---- Routes ----

@app.route('/health', methods=['GET'])
def health():
    return jsonify({"status": "ok", "coral": coral.available})


@app.route('/v1/image2text', methods=['POST'])
def image2text():
    t0 = time.time()
    request_id = str(uuid.uuid4())
    data, status = parse_image_from_request(request)
    if SHOW_GUI and cv2 is not None:
        _gui_show_message("processing request {}".format(request_id[:8]))

    if data is None:
        if status == 'too_large':
            if SHOW_GUI and cv2 is not None:
                _gui_show_message("error: payload too large")
            return jsonify({"error": "payload too large"}), 413
        if status in ('not_found', 'empty'):
            if SHOW_GUI and cv2 is not None:
                _gui_show_message("error: image missing")
            return jsonify({"error": "image is required"}), 400
        if SHOW_GUI and cv2 is not None:
            _gui_show_message("error: invalid base64")
        return jsonify({"error": "invalid image_base64"}), 400

    img_rgb = image_io.decode_image(data)
    if img_rgb is None:
        if SHOW_GUI and cv2 is not None:
            _gui_show_message("error: unsupported format")
        return jsonify({"error": "unsupported image format"}), 415

    objects = []
    coral_used = False
    try:
        objects, coral_used = coral.infer(img_rgb)
    except Exception as exc:
        print("[WARN] Coral failure {}: {}".format(request_id, exc), file=sys.stderr)
        objects = []
        coral_used = False

    boxes, east_used = txt_detector.detect(img_rgb)
    h, w = img_rgb.shape[:2]
    text_locator = "east" if east_used else "opencv_fallback"
    if not east_used and len(boxes) == 1:
        try:
            if tuple([int(v) for v in boxes[0]]) == (0, 0, w, h):
                text_locator = "full_image"
        except Exception:
            pass

    prep_options = {
        "upscale": OCR_UPSCALE,
        "clahe": OCR_CLAHE,
        "bilateral": OCR_BILATERAL,
        "blur": OCR_BLUR,
        "invert": OCR_INVERT,
        "use_threshold": OCR_USE_THRESHOLD,
    }

    ocr_lang_selected = OCR_LANG
    ocr_lang_strategy = "fixed"
    if OCR_AUTO_LANG:
        try:
            ocr_lang_selected, ocr_lang_strategy = ocr.choose_ocr_lang(
                img_rgb,
                boxes,
                requested_lang=OCR_LANG,
                prep_options=prep_options,
            )
        except Exception as exc:
            print("[WARN] OCR lang auto-detect failed {}: {}".format(request_id, exc), file=sys.stderr)
            ocr_lang_selected = OCR_LANG
            ocr_lang_strategy = "auto_error_fallback"

    text, blocks = ocr.run_ocr(
        img_rgb,
        boxes,
        lang=ocr_lang_selected,
        config=OCR_CONFIG,
        return_boxes=RETURN_BOXES,
        prep_options=prep_options,
    )
    ocr_passes = 1

    # Retry with requested language set when auto-picked language produced no text.
    if (not (text or '').strip()) and ocr_lang_selected != OCR_LANG:
        retry_text, retry_blocks = ocr.run_ocr(
            img_rgb,
            boxes,
            lang=OCR_LANG,
            config=OCR_CONFIG,
            return_boxes=RETURN_BOXES,
            prep_options=prep_options,
        )
        ocr_passes += 1
        if (retry_text or '').strip():
            text, blocks = retry_text, retry_blocks
            ocr_lang_selected = OCR_LANG
            ocr_lang_strategy = "{}+lang_retry".format(ocr_lang_strategy)

    # Retry with full-frame OCR if EAST boxes were too aggressive and OCR is empty.
    if not (text or '').strip():
        h, w = img_rgb.shape[:2]
        full_box = (0, 0, w, h)
        needs_full_pass = True
        if len(boxes) == 1:
            try:
                cur_box = tuple([int(v) for v in boxes[0]])
                needs_full_pass = cur_box != full_box
            except Exception:
                needs_full_pass = True
        if needs_full_pass:
            retry_text, retry_blocks = ocr.run_ocr(
                img_rgb,
                [full_box],
                lang=ocr_lang_selected,
                config=OCR_CONFIG,
                return_boxes=RETURN_BOXES,
                prep_options=prep_options,
            )
            ocr_passes += 1
            if (retry_text or '').strip():
                text, blocks = retry_text, retry_blocks
                ocr_lang_strategy = "{}+full_frame_retry".format(ocr_lang_strategy)

    scene_summary = schema.summarize_scene(objects)

    latency_ms = int((time.time() - t0) * 1000)

    debug_requested = DEBUG_VISUAL or request.args.get('debug', '').lower() in ('1', 'true', 'yes')
    if SHOW_GUI:
        debug_requested = True  # ensure overlay built when GUI requested
    debug_image_b64 = None
    debug_image_arr = None
    if debug_requested:
        target_w = GUI_TARGET_WIDTH if SHOW_GUI else 1024
        debug_image_b64, debug_image_arr = build_debug_image(
            img_rgb,
            boxes,
            blocks if RETURN_BOXES else [],
            objects,
            target_width=target_w,
            return_image=SHOW_GUI,
        )
        if SHOW_GUI and cv2 is not None:
            try:
                gui_frame = build_text_gui_image(
                    img_rgb,
                    boxes,
                    blocks if RETURN_BOXES else [],
                    target_width=GUI_TARGET_WIDTH,
                )
                if gui_frame is None:
                    gui_frame = debug_image_arr
                if gui_frame is not None:
                    _gui_show(gui_frame)
            except Exception:
                pass

    resp = schema.build_response(
        text=text,
        blocks=blocks,
        objects=objects,
        scene_summary=scene_summary,
        metadata={
            "latency_ms": latency_ms,
            "coral_used": coral_used,
            "east_used": east_used,
            "text_locator": text_locator,
            "ocr_lang": ocr_lang_selected,
            "ocr_lang_requested": OCR_LANG,
            "ocr_lang_strategy": ocr_lang_strategy,
            "ocr_passes": ocr_passes,
            "request_id": request_id,
            "debug": debug_requested,
        },
        return_boxes=RETURN_BOXES,
    )

    if debug_image_b64:
        resp["debug_image_base64"] = debug_image_b64

    print("[INFO] req={} size={}KB coral={} east={} latency={}ms".format(
        request_id,
        len(data) // 1024,
        coral_used,
        east_used,
        latency_ms
    ))

    return jsonify(resp)


# ---- Entrypoint ----


def main():
    parser = argparse.ArgumentParser(description='Image to text API')
    parser.add_argument('--host', default='0.0.0.0')
    parser.add_argument('--port', type=int, default=8000)
    parser.add_argument('--debug-visual', action='store_true', help='force returning debug overlay base64')
    parser.add_argument('--show-gui', action='store_true', help='show overlay in a window (requires DISPLAY and OpenCV GUI); forces debug overlay creation')
    args = parser.parse_args()

    global DEBUG_VISUAL
    if args.debug_visual:
        DEBUG_VISUAL = True
    global SHOW_GUI
    if args.show_gui:
        SHOW_GUI = True
    if SHOW_GUI:
        _start_gui_thread()
        _show_placeholder()

    app.run(host=args.host, port=args.port, threaded=True)


if __name__ == '__main__':
    main()