"""EXIF + OCR test backend — validate that GPS survives upload pipeline."""
import io, json, os, re, uuid, sqlite3, urllib.request
from pathlib import Path

from fastapi import FastAPI, File, Form, HTTPException, Query, UploadFile
from fastapi.staticfiles import StaticFiles
from PIL import Image as PILImage

# Optional HEIC/HEIF support for iPhone photos
try:
    from pillow_heif import register_heif_opener
    register_heif_opener()
    HAS_HEIF = True
except ImportError:
    HAS_HEIF = False

try:
    import pytesseract
    HAS_TESSERACT = True
except ImportError:
    HAS_TESSERACT = False

# === LLM OCR via OpenCode Go ===
OPENCODE_GO_KEY = os.environ.get("OPENCODE_GO_API_KEY", "")
OPENCODE_GO_BASE = os.environ.get("OPENCODE_GO_BASE_URL", "https://opencode.ai/zen/go/v1")
LLM_OCR_MODEL = os.environ.get("LLM_OCR_MODEL", "mimo-v2-omni")

UPLOADS = Path(__file__).parent / "uploads"
UPLOADS.mkdir(exist_ok=True)
CANTEEN_DB = Path(__file__).parent.parent / "canteen-asset-tracker" / "assets.db"

app = FastAPI(title="EXIF Test")


def _get_canteen_db():
    """Get a read/write connection to the canteen assets database."""
    if not CANTEEN_DB.exists():
        return None
    conn = sqlite3.connect(str(CANTEEN_DB))
    conn.row_factory = sqlite3.Row
    return conn


def _dms_to_decimal(dms, ref):
    """Convert EXIF DMS tuple to decimal degrees."""
    try:
        deg, minutes, sec = float(dms[0]), float(dms[1]), float(dms[2])
        decimal = deg + minutes / 60.0 + sec / 3600.0
        if ref in ("S", "W"):
            decimal = -decimal
        return round(decimal, 7)
    except Exception:
        return None


def extract_exif(image_bytes: bytes) -> dict:
    """Pull all useful EXIF fields + GPS from raw image bytes."""
    result = {"has_exif": False, "tags": {}, "gps": None}
    try:
        img = PILImage.open(io.BytesIO(image_bytes))
        exif = img.getexif()
        if not exif:
            return result

        for tag_id, value in exif.items():
            tag_name = PILImage.ExifTags.TAGS.get(tag_id, f"0x{tag_id:04x}")
            if tag_name in ("MakerNote", "UserComment", "PrintImageMatching"):
                continue
            result["tags"][tag_name] = str(value)[:300]

        if result["tags"]:
            result["has_exif"] = True

        gps_ifd = exif.get_ifd(0x8825)
        if gps_ifd:
            lat_ref = gps_ifd.get(1, "N")
            lat_dms = gps_ifd.get(2)
            lng_ref = gps_ifd.get(3, "E")
            lng_dms = gps_ifd.get(4)

            if lat_dms and lng_dms:
                lat = _dms_to_decimal(lat_dms, lat_ref)
                lng = _dms_to_decimal(lng_dms, lng_ref)
                if lat is not None and lng is not None:
                    result["gps"] = {
                        "lat": lat,
                        "lng": lng,
                        "lat_ref": str(lat_ref),
                        "lng_ref": str(lng_ref),
                        "raw_lat_dms": [float(v) for v in lat_dms],
                        "raw_lng_dms": [float(v) for v in lng_dms],
                    }
    except Exception as e:
        result["error"] = str(e)

    return result


def run_ocr(image_bytes: bytes) -> dict:
    """Run Tesseract OCR on the image."""
    if not HAS_TESSERACT:
        return {"available": False, "text": "", "error": "pytesseract not installed"}

    tmp_path = UPLOADS / f"ocr_{uuid.uuid4().hex}.jpg"
    tmp_path.write_bytes(image_bytes)
    try:
        img = PILImage.open(tmp_path)
        img_gray = img.convert("L")
        text = pytesseract.image_to_string(img_gray, config="--psm 6")
        match_5dash = re.search(r"(\d{5})[-\s]*(\d{6,})", text)
        match_5plus = re.search(r"(\d{5,})", text)
        return {
            "available": True,
            "raw_text": text.strip()[:500],
            "match_5dash6": match_5dash.group(0) if match_5dash else None,
            "match_5plus": match_5plus.group(0) if match_5plus else None,
        }
    finally:
        tmp_path.unlink(missing_ok=True)


# Maximum images per batch API call (prevents context limit issues)
BATCH_SIZE_LIMIT = 20
# Downscale images to this max dimension before sending to LLM (saves image tokens)
LLM_IMAGE_MAX_DIM = 1600


def _resize_for_llm(image_bytes: bytes, max_dim: int = LLM_IMAGE_MAX_DIM) -> bytes:
    """Downscale image to max_dim on longest edge to save vision API costs."""
    img = PILImage.open(io.BytesIO(image_bytes))
    w, h = img.size
    if max(w, h) <= max_dim:
        return image_bytes  # already small enough
    ratio = max_dim / max(w, h)
    new_size = (int(w * ratio), int(h * ratio))
    img_resized = img.resize(new_size, PILImage.LANCZOS)
    buf = io.BytesIO()
    ext = img.format or "JPEG"
    # Convert PNG, WebP, etc. to JPEG for smaller size
    if ext.upper() in ("PNG", "WEBP", "TIFF", "BMP"):
        img_resized = img_resized.convert("RGB")
        ext = "JPEG"
    img_resized.save(buf, format=ext, quality=85)
    return buf.getvalue()


def run_ocr_llm(image_bytes: bytes, model: str | None = None) -> dict:
    """Run OCR via an LLM vision model on OpenCode Go.

    Falls back to Tesseract if the API key is missing or the call fails.
    Returns the same shape as run_ocr() with an additional ''engine'' field.
    """
    if not OPENCODE_GO_KEY:
        result = run_ocr(image_bytes)
        result["engine"] = "tesseract"
        result["llm_fallback_reason"] = "no_api_key"
        return result

    import base64

    model = model or LLM_OCR_MODEL
    # Downscale to save costs
    image_bytes = _resize_for_llm(image_bytes)
    b64 = base64.b64encode(image_bytes).decode()

    body = json.dumps({
        "model": model,
        "messages": [{
            "role": "user",
            "content": [
                {"type": "text", "text": "Read ALL text and numbers visible in this photo. "
                 "Return the exact text shown, nothing else."},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}}
            ]
        }],
        "max_tokens": 200,
    }).encode()

    req = urllib.request.Request(
        f"{OPENCODE_GO_BASE}/chat/completions",
        data=body,
        headers={
            "Authorization": f"Bearer {OPENCODE_GO_KEY}",
            "Content-Type": "application/json",
            "User-Agent": "Hermes-Agent/1.0",
        },
    )

    try:
        resp = urllib.request.urlopen(req, timeout=60)
        result = json.loads(resp.read())
        text = result["choices"][0]["message"]["content"].strip()
    except Exception as exc:
        # Fallback to Tesseract
        result = run_ocr(image_bytes)
        result["engine"] = "tesseract"
        result["llm_fallback_reason"] = str(exc)[:200]
        return result

    match_5dash = re.search(r"(\d{5})[-\s]*(\d{6,})", text)
    match_5plus = re.search(r"(\d{5,})", text)

    return {
        "available": True,
        "engine": "llm",
        "llm_model": model,
        "raw_text": text[:500],
        "match_5dash6": match_5dash.group(0) if match_5dash else None,
        "match_5plus": match_5plus.group(0) if match_5plus else None,
    }


def run_ocr_llm_batch(images: list[tuple[str, bytes]], model: str | None = None, max_per_batch: int = BATCH_SIZE_LIMIT) -> list[dict]:
    """Batch OCR multiple images in a single LLM API call.

    Sends all images in one request with a structured JSON prompt.
    Returns list of OCR results in the same order as the input images.
    Falls back to individual calls if the batch call fails.

    images: list of (filename, image_bytes)
    """
    import base64

    if not OPENCODE_GO_KEY:
        return [run_ocr(b) for _, b in images]

    model = model or LLM_OCR_MODEL

    # Split into sub-batches if over the limit
    all_results: list[dict] = []
    for batch_start in range(0, len(images), max_per_batch):
        batch = images[batch_start:batch_start + max_per_batch]
        batch_results = _run_ocr_llm_batch_inner(batch, model)
        all_results.extend(batch_results)
    return all_results


def _run_ocr_llm_batch_inner(batch: list[tuple[str, bytes]], model: str) -> list[dict]:
    """Inner helper: sends one batch of images in a single API call."""
    import base64

    # Resize all images first
    resized = [(_resize_for_llm(b), n) for n, b in batch]

    content: list[dict] = [{
        "type": "text",
        "text": (
            f"I have {len(batch)} photos. For EACH photo, read ALL visible text and numbers.\n"
            "Respond with a JSON array of objects, one per photo in order:\n"
            '[{"i":0,"text":"all text and digits found","digits":"e.g. 12345-678901 or null if none"}, ...]\n'
            'Return ONLY the JSON array, no markdown, no explanation.'
        )
    }]

    for idx, (img_bytes, _) in enumerate(resized):
        b64 = base64.b64encode(img_bytes).decode()
        content.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{b64}"}
        })

    body = json.dumps({
        "model": model,
        "messages": [{"role": "user", "content": content}],
        "max_tokens": min(300 * len(batch), 8000),
        "temperature": 0.1,
    }).encode()

    req = urllib.request.Request(
        f"{OPENCODE_GO_BASE}/chat/completions",
        data=body,
        headers={
            "Authorization": f"Bearer {OPENCODE_GO_KEY}",
            "Content-Type": "application/json",
            "User-Agent": "Hermes-Agent/1.0",
        },
    )

    try:
        resp = urllib.request.urlopen(req, timeout=120)
        result = json.loads(resp.read())
        raw = result["choices"][0]["message"]["content"].strip()
    except Exception:
        # Fallback: individual calls for this batch
        return [run_ocr_llm(b, model) for n, b in batch]

    # Strip markdown code fences if present
    raw_clean = raw
    if raw_clean.startswith("```"):
        raw_clean = raw_clean.split("\n", 1)[-1]
        raw_clean = raw_clean.rsplit("```", 1)[0].strip()

    try:
        parsed = json.loads(raw_clean)
        results: list[dict] = []
        for item in parsed:
            raw_text = str(item.get("text", "") or "")
            digit_str = str(item.get("digits") or "")
            combined = raw_text + " " + digit_str
            match_5dash = re.search(r"(\d{5})[-\s]*(\d{6,})", combined)
            match_5plus = re.search(r"(\d{5,})", combined)
            results.append({
                "available": True,
                "engine": "llm_batch",
                "llm_model": model,
                "raw_text": raw_text[:500],
                "match_5dash6": match_5dash.group(0) if match_5dash else None,
                "match_5plus": match_5plus.group(0) if match_5plus else None,
            })
        return results
    except (json.JSONDecodeError, KeyError, TypeError):
        # Fallback to individual calls
        return [run_ocr_llm(b, model) for n, b in batch]


def lookup_machine_id(machine_id: str) -> dict | None:
    """Look up an asset by machine_id. Returns asset dict or None."""
    conn = _get_canteen_db()
    if not conn:
        return None
    try:
        row = conn.execute(
            "SELECT id, machine_id, name, category, status, address, building_name, "
            "floor, room, latitude, longitude, make, model, description, photo_path "
            "FROM assets WHERE machine_id = ?",
            (machine_id.strip(),),
        ).fetchone()
        if row:
            return {
                "id": row["id"],
                "machine_id": row["machine_id"],
                "name": row["name"],
                "category": row["category"],
                "status": row["status"],
                "address": row["address"],
                "building_name": row["building_name"],
                "floor": row["floor"],
                "room": row["room"],
                "latitude": row["latitude"],
                "longitude": row["longitude"],
                "make": row["make"],
                "model": row["model"],
                "description": row["description"],
                "photo_path": row["photo_path"],
            }
    finally:
        conn.close()
    return None


@app.post("/api/analyze")
async def analyze_photo(
    file: UploadFile = File(...),
    ocr_engine: str = Query(default="tesseract"),
    ocr_model: str = Query(default=""),
):
    """Upload a photo, get back EXIF + OCR results.

    Query params:
      - ocr_engine: ''tesseract'' (default) or ''llm''
      - ocr_model:  model name override (e.g. ''mimo-v2-omni'', ''glm-5.1'')
    """
    contents = await file.read()
    file_size = len(contents)

    ext = Path(file.filename or "photo.jpg").suffix.lower()
    if ext not in {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff", ".tif", ".dng", ".heic", ".heif"}:
        ext = ".jpg"
    fname = f"{uuid.uuid4().hex}{ext}"
    (UPLOADS / fname).write_bytes(contents)

    exif_result = extract_exif(contents)

    if ocr_engine == "llm":
        ocr_result = run_ocr_llm(contents, ocr_model or None)
    elif HAS_TESSERACT:
        ocr_result = run_ocr(contents)
        ocr_result["engine"] = "tesseract"
    else:
        ocr_result = {"available": False, "text": "", "engine": "none"}

    machine_id = None
    asset = None
    if ocr_result.get("match_5dash6"):
        machine_id = re.sub(r"\D", "", ocr_result["match_5dash6"])[-5:]
        asset = lookup_machine_id(machine_id)

    return {
        "filename": file.filename,
        "saved_as": fname,
        "file_size": file_size,
        "file_size_kb": round(file_size / 1024, 1),
        "exif": exif_result,
        "ocr": ocr_result,
        "machine_id": machine_id,
        "asset": asset,
    }


@app.post("/api/bulk-process")
async def bulk_process(
    files: list[UploadFile] = File(...),
    ocr_engine: str = Query(default="tesseract"),
    ocr_model: str = Query(default=""),
):
    """Process multiple photos: OCR each, extract EXIF GPS, look up matching assets.

    Query params:
      - ocr_engine: ''tesseract'' (default) or ''llm''
      - ocr_model:  model name override

    Returns a list of results, each with:
      - filename, exif (gps), ocr match, matched asset (if found)
      - needs_gps: true if asset exists AND has no coordinates AND photo has GPS
    """
    if not files:
        return {"results": [], "summary": {"total": 0, "has_gps": 0, "matched": 0, "needs_gps": 0}}

    # Read all files first
    all_files: list[tuple[str, bytes]] = []
    for file in files:
        contents = await file.read()
        all_files.append((file.filename or "photo.jpg", contents))

    ocr_use_llm = ocr_engine == "llm"

    # Batch OCR if using LLM mode
    if ocr_use_llm:
        llm_ocr_results = run_ocr_llm_batch(all_files, ocr_model or None)
    else:
        llm_ocr_results = None

    results = []
    summary = {"total": len(all_files), "has_gps": 0, "matched": 0, "needs_gps": 0}

    for i, (orig_fname, contents) in enumerate(all_files):
        file_size = len(contents)

        # Save
        ext = Path(orig_fname or "photo.jpg").suffix.lower()
        if ext not in {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff", ".tif", ".dng", ".heic", ".heif"}:
            ext = ".jpg"
        saved_name = f"{uuid.uuid4().hex}{ext}"
        (UPLOADS / saved_name).write_bytes(contents)

        exif_result = extract_exif(contents)

        if ocr_use_llm:
            ocr_result = llm_ocr_results[i] if i < len(llm_ocr_results) else {"available": False, "engine": "none", "error": "missing batch result"}
        elif HAS_TESSERACT:
            ocr_result = run_ocr(contents)
            ocr_result["engine"] = "tesseract"
        else:
            ocr_result = {"available": False, "text": "", "engine": "none"}

        has_gps = exif_result.get("gps") is not None
        if has_gps:
            summary["has_gps"] += 1

        machine_id = None
        asset = None
        needs_gps = False

        if ocr_result.get("match_5dash6"):
            machine_id = re.sub(r"\D", "", ocr_result["match_5dash6"])[-5:]
            asset = lookup_machine_id(machine_id)
            if asset:
                summary["matched"] += 1
                # Check if asset needs GPS
                if (asset["latitude"] is None or asset["longitude"] is None) and has_gps:
                    needs_gps = True
                    summary["needs_gps"] += 1

        results.append({
            "filename": orig_fname,
            "saved_as": saved_name,
            "file_size_kb": round(file_size / 1024, 1),
            "exif": exif_result,
            "ocr": ocr_result,
            "machine_id": machine_id,
            "asset": asset,
            "needs_gps": needs_gps,
        })

    return {"results": results, "summary": summary}


@app.get("/api/lookup")
async def lookup_asset(machine_id: str = ""):
    """Look up an asset by machine_id in the canteen assets database."""
    if not machine_id or not machine_id.strip():
        return {"found": False, "reason": "No machine_id provided"}

    asset = lookup_machine_id(machine_id.strip())
    if asset:
        return {"found": True, "asset": asset}
    return {"found": False, "reason": f"No asset with machine_id '{machine_id}'"}


@app.post("/api/push-gps")
async def push_gps(request: dict):
    """Update an asset's GPS coordinates from photo EXIF data.

    Body: {"asset_id": 123, "latitude": 40.7417, "longitude": -73.9292}
    Only updates assets that currently have NULL lat/lng.
    """
    asset_id = request.get("asset_id")
    latitude = request.get("latitude")
    longitude = request.get("longitude")

    if not asset_id or latitude is None or longitude is None:
        raise HTTPException(400, "asset_id, latitude, and longitude are required")

    conn = _get_canteen_db()
    if not conn:
        raise HTTPException(500, "Database not available")

    try:
        # Only update if currently NULL
        row = conn.execute(
            "SELECT latitude, longitude FROM assets WHERE id = ?", (int(asset_id),)
        ).fetchone()
        if not row:
            conn.close()
            raise HTTPException(404, f"Asset {asset_id} not found")
        if row["latitude"] is not None and row["longitude"] is not None:
            conn.close()
            return {"updated": False, "reason": "Asset already has GPS coordinates", "asset_id": asset_id}

        conn.execute(
            "UPDATE assets SET latitude = ?, longitude = ?, updated_at = datetime('now') WHERE id = ?",
            (float(latitude), float(longitude), int(asset_id)),
        )
        conn.commit()
        conn.close()
        return {"updated": True, "asset_id": asset_id, "latitude": float(latitude), "longitude": float(longitude)}
    except Exception as e:
        try:
            conn.close()
        except Exception:
            pass
        raise HTTPException(500, str(e))


@app.get("/api/uploads")
async def list_uploads():
    """List previously uploaded files."""
    files = sorted(UPLOADS.glob("*"), key=lambda p: p.stat().st_mtime, reverse=True)
    return [
        {"name": f.name, "size_kb": round(f.stat().st_size / 1024, 1)}
        for f in files[:20]
    ]


# Mount static LAST
app.mount("/", StaticFiles(directory="static", html=True), name="static")

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