import os
import sys
import time
from pathlib import Path

import numpy as np
import streamlit as st
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parent))

try:
    from .ocr_engine import DEFAULT_MODEL_ID, OCREngine
    from .storage import is_model_cached, storage_info
    from .utils import clean_ocr_markdown
except ImportError:
    from ocr_engine import DEFAULT_MODEL_ID, OCREngine
    from storage import is_model_cached, storage_info
    from utils import clean_ocr_markdown


st.set_page_config(
    page_title="LightOnOCR Handwriting POC",
    page_icon="📝",
    layout="wide",
    initial_sidebar_state="expanded",
)

st.markdown(
    """
    <style>
    .stApp {
        background: #f7f8fb;
        color: #111827;
    }
    [data-testid="stSidebar"] {
        background: #ffffff;
        border-right: 1px solid #e5e7eb;
    }
    .hero {
        padding: 1.35rem 0 1rem;
        border-bottom: 1px solid #e5e7eb;
        margin-bottom: 1rem;
    }
    .hero h1 {
        font-size: 2rem;
        line-height: 1.12;
        margin: 0 0 .35rem;
        color: #0f172a;
    }
    .hero p {
        margin: 0;
        color: #475569;
        max-width: 780px;
    }
    .metric-strip {
        display: grid;
        grid-template-columns: repeat(3, minmax(0, 1fr));
        gap: .75rem;
        margin: .75rem 0 1rem;
        max-width: 100%;
        overflow: hidden;
    }
    .metric-box {
        background: #ffffff;
        border: 1px solid #e5e7eb;
        border-radius: 8px;
        padding: .8rem .9rem;
        min-width: 0;
        overflow: hidden;
    }
    .metric-box span {
        display: block;
        color: #64748b;
        font-size: .78rem;
    }
    .metric-box strong {
        color: #0f172a;
        font-size: .95rem;
        overflow-wrap: anywhere;
        word-break: break-word;
    }
    [data-testid="column"] {
        min-width: 0 !important;
    }
    [data-testid="stImage"] img {
        max-height: 72vh;
        object-fit: contain;
    }
    [data-testid="stCodeBlock"] {
        max-width: 100%;
        overflow-x: auto;
    }
    .result-empty {
        align-items: center;
        color: #64748b;
        display: flex;
        min-height: 260px;
    }
    .stButton > button {
        border-radius: 8px;
        font-weight: 700;
        min-height: 2.75rem;
    }
    div[data-testid="stFileUploader"] {
        background: #ffffff;
        border: 1px solid #e5e7eb;
        border-radius: 8px;
        padding: .75rem;
    }
    @media (max-width: 800px) {
        .metric-strip { grid-template-columns: 1fr; }
        .hero h1 { font-size: 1.55rem; }
    }
    </style>
    """,
    unsafe_allow_html=True,
)

st.markdown(
    """
    <div class="hero">
      <h1>LightOnOCR Handwriting POC</h1>
      <p>Ekstraksi teks dokumen, nota, dan handwriting dengan model PetaniHandal berbasis LightOnOCR. Output dikembalikan sebagai Markdown agar mudah diaudit, disalin, atau diproses lanjut.</p>
    </div>
    """,
    unsafe_allow_html=True,
)


PRESET_LABELS = {
    "handwriting": "Handwriting",
    "document": "Dokumen umum",
    "receipt": "Nota / invoice",
}


with st.sidebar:
    st.header("Pengaturan POC")
    preset = st.segmented_control(
        "Mode ekstraksi",
        options=list(PRESET_LABELS.keys()),
        format_func=lambda key: PRESET_LABELS[key],
        default="handwriting",
    )
    max_size = st.slider(
        "Resolusi sisi terpanjang",
        960,
        2200,
        1540,
        step=100,
        help="Rekomendasi LightOnOCR sekitar 1540px untuk menjaga geometri teks.",
    )
    max_new_tokens = st.slider("Batas token output", 512, 8192, 4096, step=512)
    temperature = st.slider("Temperature", 0.0, 0.7, 0.1, step=0.05)

    st.divider()
    st.caption("Model aktif")
    st.code(DEFAULT_MODEL_ID, language=None)

    info = storage_info()
    cached = is_model_cached(DEFAULT_MODEL_ID)
    with st.expander("Storage & cache", expanded=False):
        st.write(f"Environment: {info['environment']}")
        st.write(f"Base path: `{info['base_path']}`")
        st.write("Model cache:", "tersedia" if cached else "belum tersedia")
        if "disk" in info:
            disk = info["disk"]
            used_pct = disk["used_gb"] / disk["total_gb"] if disk["total_gb"] else 0
            st.progress(used_pct, text=f"{disk['used_gb']} GB / {disk['total_gb']} GB digunakan")


@st.cache_resource(show_spinner="Menyiapkan LightOnOCR. Download hanya terjadi jika model belum ada di cache.")
def get_ocr_engine(model_preset: str, output_tokens: int, temp: float):
    return OCREngine(
        preset=model_preset,
        max_new_tokens=output_tokens,
        temperature=temp,
    )


runtime_label = "vLLM endpoint" if os.getenv("LIGHTONOCR_ENDPOINT_URL") else "Transformers local"
model_id = os.getenv("LIGHTONOCR_MODEL_ID", DEFAULT_MODEL_ID)
cache_label = "Warm cache" if is_model_cached(model_id) else "Cache pending"
st.markdown(
    f"""
    <div class="metric-strip">
      <div class="metric-box"><span>Runtime</span><strong>{runtime_label}</strong></div>
      <div class="metric-box"><span>Model</span><strong>{model_id}</strong></div>
      <div class="metric-box"><span>Cache</span><strong>{cache_label}</strong></div>
    </div>
    """,
    unsafe_allow_html=True,
)

uploaded_file = st.file_uploader(
    "Unggah gambar dokumen",
    type=["png", "jpg", "jpeg"],
    accept_multiple_files=False,
)

if "ocr_result" not in st.session_state:
    st.session_state.ocr_result = None
if "ocr_elapsed" not in st.session_state:
    st.session_state.ocr_elapsed = None
if "uploaded_name" not in st.session_state:
    st.session_state.uploaded_name = None

if uploaded_file is None:
    st.session_state.uploaded_name = None
    st.info("Unggah satu gambar untuk memulai ekstraksi.")
else:
    if st.session_state.uploaded_name != uploaded_file.name:
        st.session_state.uploaded_name = uploaded_file.name
        st.session_state.ocr_result = None
        st.session_state.ocr_elapsed = None

    image = Image.open(uploaded_file).convert("RGB")
    img_array_bgr = np.array(image)[:, :, ::-1].copy()
    h, w = img_array_bgr.shape[:2]

    left, right = st.columns([1, 1], gap="medium")
    with left:
        st.subheader("Preview Dokumen")
        st.image(image, use_container_width=True)
        st.caption(f"Ukuran asli: {w} x {h}px")

    with right:
        st.subheader("Hasil OCR")
        run = st.button("Jalankan Ekstraksi", type="primary", use_container_width=True)
        if run:
            with st.spinner("Membaca dokumen dan menyusun Markdown..."):
                engine = get_ocr_engine(preset, max_new_tokens, temperature)
                t0 = time.time()
                result = engine.process_image(img_array_bgr, max_size=max_size)
                elapsed = time.time() - t0

            md_text = clean_ocr_markdown(result.get("markdown_text", ""))
            st.session_state.ocr_result = md_text
            st.session_state.ocr_elapsed = elapsed

        if st.session_state.ocr_elapsed is not None:
            st.success(f"Selesai dalam {st.session_state.ocr_elapsed:.1f} detik")

        with st.container(border=True):
            if st.session_state.ocr_result:
                tab_rendered, tab_raw = st.tabs(["Rendered", "Markdown"])
                with tab_rendered:
                    st.markdown(st.session_state.ocr_result, unsafe_allow_html=True)
                with tab_raw:
                    st.code(st.session_state.ocr_result, language="markdown")
            elif st.session_state.ocr_elapsed is not None:
                st.warning("Model tidak mengembalikan teks untuk gambar ini.")
            else:
                st.markdown(
                    '<div class="result-empty">Hasil OCR akan muncul di sini setelah ekstraksi dijalankan.</div>',
                    unsafe_allow_html=True,
                )