import streamlit as st
import os
import json
import datetime
import asyncio
import threading

from VerifierApp.config import INPUTS_DIR, MODELS_DIR, RUNS_DIR, get_folders_in_dir, get_files_in_dir, get_base_prompt, \
    save_base_prompt
from VerifierApp.core.models import fetch_available_models
from VerifierApp.core.extractor import run_extraction
from VerifierApp.core.verifier import verify_phenotypes
from VerifierApp.core import job_manager
from VerifierApp.localisation_module import t, ICONS

_STATUS_LABELS = {
    job_manager.STATUS_PENDING:   "testing_bg_status_pending",
    job_manager.STATUS_RUNNING:   "testing_bg_status_running",
    job_manager.STATUS_DONE:      "testing_bg_status_done",
    job_manager.STATUS_FAILED:    "testing_bg_status_failed",
    job_manager.STATUS_CANCELLED: "testing_bg_status_cancelled",
}

def _job_fn(
    input_group_path: str,
    model_group_path: str,
    selected_llm: str,
    custom_prompt: str,
    selected_input_group: str,
    selected_model_group: str,
    cancel_event: threading.Event,
    repetition: int = 1,
    milestone: int = 1,
    use_model_context: bool = True,
) -> dict:
    """Single extraction + verification cycle, cancellation-aware."""

    files_in_model = get_files_in_dir(model_group_path)

    file_for_llm = next((f for f in files_in_model if f.endswith('.bnet')), None) or \
                   next((f for f in files_in_model if f.endswith('.aeon')), None)
    file_for_aeon = next((f for f in files_in_model if f.endswith('.aeon')), None)

    if not file_for_llm or not file_for_aeon:
        raise ValueError(t("testing_error_no_model_files"))

    path_for_llm = os.path.join(model_group_path, file_for_llm)
    path_for_aeon = os.path.join(model_group_path, file_for_aeon)

    stamp = datetime.datetime.now().strftime('%H%M%S%f')
    temp_phenotypes_path = os.path.join(RUNS_DIR, f"temp_phenotypes_{stamp}.json")
    temp_report_path = os.path.join(RUNS_DIR, f"temp_report_{stamp}.json")

    try:
        if cancel_event.is_set():
            raise RuntimeError(t("testing_cancelled_before_extraction"))

        extracted_json = asyncio.run(run_extraction(
            input_path_str=input_group_path,
            output_path_str=temp_phenotypes_path,
            model_path_str=path_for_llm if use_model_context else None,
            custom_prompt=custom_prompt,
            llm_model_name=selected_llm,
        ))

        if not extracted_json:
            raise ValueError(t("testing_error_extraction_failed"))

        extraction_usage = extracted_json.get("usage", {}) if isinstance(extracted_json, dict) else {}

        try:
            with open(temp_phenotypes_path, 'r', encoding='utf-8') as _f:
                _raw = _f.read().strip()
            if not _raw:
                raise ValueError(t("testing_error_extraction_failed"))
            _parsed = json.loads(_raw)
            if isinstance(_parsed, str):
                _parsed = json.loads(_parsed)
            if not isinstance(_parsed, list):
                raise ValueError(t("testing_error_extraction_failed"))
        except json.JSONDecodeError as _e:
            raise ValueError(f"{t('testing_error_extraction_failed')} ({_e})")

        if cancel_event.is_set():
            raise RuntimeError(t("testing_cancelled_before_verification"))

        verify_phenotypes(
            model_path=path_for_aeon,
            json_path=temp_phenotypes_path,
            report_path=temp_report_path,
        )

        with open(temp_report_path, 'r', encoding='utf-8') as f:
            final_report = json.load(f)

        timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        run_id = f"run_{timestamp}_{selected_llm.replace(':', '_').replace('/', '_')}_run_{repetition}"
        run_data = {
            "run_id":      run_id,
            "timestamp":   timestamp,
            "milestone":   milestone,
            "input_group": selected_input_group,
            "model_group": selected_model_group,
            "llm_used":    selected_llm,
            "prompt":      custom_prompt,
            "tokens_used": extraction_usage,
            "results":     final_report,
        }
        run_filename = f"{run_id}.json"
        with open(os.path.join(RUNS_DIR, run_filename), "w", encoding="utf-8") as f:
            json.dump(run_data, f, indent=4, ensure_ascii=False)

        try:
            from VerifierApp.core.runs_db import upsert_run
            upsert_run(RUNS_DIR, run_filename, run_data)
        except Exception:
            pass   # DB write failure must never break a completed run

        return {"success": True, "report": final_report, "run_id": run_id, "usage": extraction_usage}

    finally:
        for p in (temp_phenotypes_path, temp_report_path):
            if os.path.exists(p):
                try:
                    os.remove(p)
                except OSError:
                    pass


def _submit_jobs(
    input_group_path, model_group_path, selected_llm,
    custom_prompt, selected_input_group, selected_model_group,
    repetitions: int = 1,
    group_id=None,
    milestone: int = 1,
    use_model_context: bool = True,
) -> list:
    """Submit *repetitions* background jobs and return their IDs."""
    ids = []
    for rep in range(1, repetitions + 1):
        label = (
            f"{selected_llm} | {selected_input_group} → {selected_model_group}"
            + (f" ({t('testing_bg_job_rep', rep=rep, total=repetitions)})" if repetitions > 1 else "")
        )
        jid = job_manager.submit_job(
            label=label,
            fn=_job_fn,
            kwargs=dict(
                input_group_path=input_group_path,
                model_group_path=model_group_path,
                selected_llm=selected_llm,
                custom_prompt=custom_prompt,
                selected_input_group=selected_input_group,
                selected_model_group=selected_model_group,
                repetition=rep,
                milestone=milestone,
                use_model_context=use_model_context,
            ),
            group_id=group_id,
        )
        ids.append(jid)
    return ids


_ACTIVE_STATUSES = frozenset({job_manager.STATUS_PENDING, job_manager.STATUS_RUNNING})
_FINISHED_STATUSES = frozenset({job_manager.STATUS_DONE, job_manager.STATUS_FAILED, job_manager.STATUS_CANCELLED})


def _draw_job_monitor_content():
    """Shared rendering logic for the job monitor (called from both fragment variants)."""
    all_jobs = job_manager.get_all_jobs()

    if not all_jobs:
        st.caption(t("testing_bg_jobs_none"))
        return

    active_jobs   = [j for j in all_jobs if j["status"] in _ACTIVE_STATUSES]
    finished_jobs = [j for j in all_jobs if j["status"] in _FINISHED_STATUSES]

    col_hdr, col_clear = st.columns([3, 1])
    with col_hdr:
        active_count = len(active_jobs)
        st.caption(t("testing_bg_active", n=active_count) if active_count else t("testing_bg_all_done"))
    with col_clear:
        if finished_jobs and st.button(t("testing_bg_jobs_clear"), key="clear_finished", use_container_width=True):
            job_manager.clear_finished()

    shown_jobs = active_jobs + finished_jobs[:10]
    shown_jobs.sort(key=lambda j: (j["status"] in _FINISHED_STATUSES, j["created_at"]))

    for job in shown_jobs:
        status_label = t(_STATUS_LABELS.get(job["status"], job["status"]))
        with st.container(border=True):
            col_info, col_act = st.columns([5, 1])
            with col_info:
                st.write(f"**{status_label}** — {job['label']}")
                st.caption(
                    f"ID: `{job['job_id']}` | Started: {job['created_at']}"
                    + (f" | Finished: {job['finished_at']}" if job["finished_at"] else "")
                )
                if job["status"] == job_manager.STATUS_FAILED and job.get("error"):
                    st.error(job["error"])
                elif job["status"] == job_manager.STATUS_DONE and job.get("result"):
                    res = job["result"]
                    report = res.get("report", {})
                    results = report.get("results", []) if isinstance(report, dict) else (report if isinstance(report, list) else [])
                    verified = sum(1 for r in results if isinstance(r, dict) and r.get("verified"))
                    usage = res.get("usage", {})
                    token_info = f" | {ICONS['success']} {usage.get('total_tokens', '?')} tokens" if usage else ""
                    st.success(f"{verified}/{len(results)} verified — saved as `{res.get('run_id', '')}`{token_info}")
            with col_act:
                if job["status"] in _ACTIVE_STATUSES:
                    if st.button(t("testing_bg_jobs_cancel"), key=f"cancel_{job['job_id']}", use_container_width=True):
                        job_manager.cancel_job(job["job_id"])

    hidden = len(finished_jobs) - 10
    if hidden > 0:
        st.caption(f"… and {hidden} more finished jobs (use 'Clear finished' to remove them)")


@st.fragment(run_every="2s")
def _render_job_monitor_polling():
    """Auto-refreshing monitor — only mounted when jobs are active.
    Triggers a full-page rerun once all jobs finish so the parent can
    switch to the idle (non-polling) variant."""
    _draw_job_monitor_content()
    if job_manager.get_active_count() == 0:
        st.rerun()


@st.fragment
def _render_job_monitor_idle():
    """Static monitor — no periodic polling.  Used when no jobs are running."""
    _draw_job_monitor_content()



def render_testing():
    st.header(t("testing_header"))

    available_input_dirs = get_folders_in_dir(INPUTS_DIR)
    available_model_dirs = get_folders_in_dir(MODELS_DIR)

    if not available_input_dirs or not available_model_dirs:
        st.warning(f"{ICONS['warning']} {t('testing_no_data_warning')}")
        return

    # ── Background jobs panel ─────────────────────────────────────────────────
    active_count = job_manager.get_active_count()
    with st.expander(
        f"{ICONS['jobs']} {t('testing_bg_jobs_header')} — active: {active_count}",
        expanded=active_count > 0,
    ):
        if active_count > 0:
            _render_job_monitor_polling()
        else:
            _render_job_monitor_idle()

    st.markdown("---")

    tab1, tab2 = st.tabs([
        f"{ICONS['nav_testing']} {t('testing_tab_single')}",
        f":material/rocket_launch: {t('testing_tab_batch')}",
    ])

    # ── Single run tab ────────────────────────────────────────────────────────
    with tab1:
        col1, col2 = st.columns(2)
        with col1:
            st.subheader(t("testing_data_pairing"))
            selected_input_group = st.selectbox(t("testing_select_input"), available_input_dirs, key="single_in")
            selected_model_group = st.selectbox(t("testing_select_model"), available_model_dirs, key="single_mod")

        with col2:
            st.subheader(t("testing_llm_setup"))
            available_llms = fetch_available_models()
            selected_llm = st.selectbox(t("testing_select_llm"), available_llms, key="single_llm")
            repetitions = st.number_input(
                t("testing_repetitions_label"),
                min_value=1, max_value=1000, value=1, step=1,
                help=t("testing_repetitions_help"),
                key="single_reps",
            )
            single_milestone = st.number_input(
                t("testing_milestone_label"),
                min_value=1, max_value=100, value=1, step=1,
                help=t("testing_milestone_help"),
                key="single_milestone",
            )
            single_use_model_context = st.toggle(
                t("testing_use_model_context_label"),
                value=True,
                help=t("testing_use_model_context_help"),
                key="single_model_context",
            )

        st.subheader(t("testing_prompt_section"))
        current_base_prompt = get_base_prompt()
        custom_prompt = st.text_area(t("testing_prompt_label"), value=current_base_prompt, height=150, key="single_prompt")

        col_save, col_run = st.columns([1, 3])
        with col_save:
            if st.button(t('testing_save_prompt'), icon=ICONS['save'], key="single_save", use_container_width=True):
                save_base_prompt(custom_prompt)
                st.success(t("testing_prompt_saved"))
        with col_run:
            if st.button(t('testing_run_btn'), icon=ICONS['run'], type="primary", use_container_width=True, key="single_run"):
                group_id = f"single_{datetime.datetime.now().strftime('%H%M%S%f')}"
                ids = _submit_jobs(
                    os.path.join(INPUTS_DIR, selected_input_group),
                    os.path.join(MODELS_DIR, selected_model_group),
                    selected_llm, custom_prompt,
                    selected_input_group, selected_model_group,
                    repetitions=int(repetitions),
                    group_id=group_id,
                    milestone=int(single_milestone),
                    use_model_context=single_use_model_context,
                )
                st.success(t("testing_submitted_bg", n=len(ids)))
                st.rerun()

    # ── Batch run tab ─────────────────────────────────────────────────────────
    with tab2:
        st.info(t("testing_batch_info"))

        st.subheader(t("testing_batch_select_llms"))
        available_llms_b = fetch_available_models()
        batch_llms = st.multiselect(
            t("testing_batch_select_llm_label"), available_llms_b,
            default=["qwen3.5-122b"] if "qwen3.5-122b" in available_llms_b else None,
        )

        st.subheader(t("testing_batch_map"))

        if "batch_pairs" not in st.session_state:
            st.session_state.batch_pairs = [{"input": available_input_dirs[0], "model": available_model_dirs[0]}]

        for i, pair in enumerate(st.session_state.batch_pairs):
            col_in, col_mod, col_del = st.columns([4, 4, 1])
            with col_in:
                pair["input"] = st.selectbox(
                    t("testing_batch_input_label", n=i + 1), available_input_dirs,
                    index=available_input_dirs.index(pair["input"]) if pair["input"] in available_input_dirs else 0,
                    key=f"b_in_{i}",
                )
            with col_mod:
                pair["model"] = st.selectbox(
                    t("testing_batch_model_label", n=i + 1), available_model_dirs,
                    index=available_model_dirs.index(pair["model"]) if pair["model"] in available_model_dirs else 0,
                    key=f"b_mod_{i}",
                )
            with col_del:
                st.write("")
                if st.button("", icon=ICONS["delete"], key=f"b_del_{i}",
                             disabled=len(st.session_state.batch_pairs) == 1):
                    st.session_state.batch_pairs.pop(i)
                    st.rerun()

        if st.button(t("testing_batch_add_pair")):
            st.session_state.batch_pairs.append({"input": available_input_dirs[0], "model": available_model_dirs[0]})
            st.rerun()

        st.subheader(t("testing_batch_prompt_section"))
        batch_prompt = st.text_area(
            t("testing_batch_prompt_label"), value=get_base_prompt(), height=150, key="batch_prompt",
        )

        batch_reps = st.number_input(
            t("testing_repetitions_label"),
            min_value=1, max_value=1000, value=1, step=1,
            help=t("testing_repetitions_help"),
            key="batch_reps",
        )

        batch_milestone = st.number_input(
            t("testing_milestone_label"),
            min_value=1, max_value=100, value=1, step=1,
            help=t("testing_milestone_help"),
            key="batch_milestone",
        )

        batch_use_model_context = st.toggle(
            t("testing_use_model_context_label"),
            value=True,
            help=t("testing_use_model_context_help"),
            key="batch_model_context",
        )

        total_runs = len(st.session_state.batch_pairs) * len(batch_llms) * int(batch_reps)
        st.write(t("testing_batch_total_with_reps",
                   total=total_runs,
                   pairs=len(st.session_state.batch_pairs),
                   llms=len(batch_llms),
                   reps=int(batch_reps)))

        if st.button(t("testing_batch_run_btn"), type="primary", use_container_width=True,
                     disabled=(total_runs == 0)):
            group_id = f"batch_{datetime.datetime.now().strftime('%H%M%S%f')}"
            submitted = 0
            for llm in batch_llms:
                for pair in st.session_state.batch_pairs:
                    ids = _submit_jobs(
                        os.path.join(INPUTS_DIR, pair["input"]),
                        os.path.join(MODELS_DIR, pair["model"]),
                        llm, batch_prompt,
                        pair["input"], pair["model"],
                        repetitions=int(batch_reps),
                        group_id=group_id,
                        milestone=int(batch_milestone),
                        use_model_context=batch_use_model_context,
                    )
                    submitted += len(ids)
            st.success(t("testing_submitted_bg", n=submitted))
            st.rerun()
