from __future__ import annotations

import functools
import html
import importlib.util
import os
import socket
import sys
from contextlib import closing
from pathlib import Path
from typing import Any

import gradio as gr
import nltk
import psutil
import torch
from transformers import AutoModel

from example_inputs import EXAMPLE_INPUTS


def _ensure_nltk_punkt_resources() -> None:
    """Download punkt resources on first run so inference never fails."""

    for resource in ("punkt", "punkt_tab"):
        try:
            nltk.data.find(f"tokenizers/{resource}")
        except LookupError:
            nltk.download(resource, quiet=True)


DEFAULT_MODEL = "hotchpotch/open-provence-reranker-xsmall-v1"
DEFAULT_THRESHOLD = 0.1
DEFAULT_BATCH_SIZE = 8
MODEL_SCRIPT_ENV = "OPEN_PROVENCE_MODELING_SCRIPT_PATH"


DEVICE_HINT = os.getenv("OPEN_PROVENCE_INFER_DEVICE")


def _resolve_device() -> torch.device:
    if DEVICE_HINT:
        try:
            return torch.device(DEVICE_HINT)
        except (TypeError, RuntimeError):
            pass
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():  # type: ignore[attr-defined]
        return torch.device("mps")
    return torch.device("cpu")


INFERENCE_DEVICE = _resolve_device()

_ensure_nltk_punkt_resources()


def _discover_local_ips() -> list[str]:
    """Return IPv4 addresses reachable from the local network."""

    addresses: set[str] = set()
    try:
        interfaces = psutil.net_if_addrs()
    except Exception:  # pragma: no cover - platform specific
        interfaces = {}

    for entries in interfaces.values():
        for entry in entries:
            if entry.family == socket.AF_INET:
                ip = entry.address
                if ip and not ip.startswith("127."):
                    addresses.add(ip)

    if not addresses:
        try:
            fallback_ip = socket.gethostbyname(socket.gethostname())
        except OSError:
            fallback_ip = None
        if fallback_ip and not fallback_ip.startswith("127."):
            addresses.add(fallback_ip)

    return sorted(addresses)


def _find_free_port() -> int:
    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
        sock.bind(("", 0))
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        return int(sock.getsockname()[1])


def _resolve_server_port(preferred: int) -> tuple[int, bool]:
    """Return an available port, optionally falling back from the preferred value."""

    if preferred <= 0:
        return _find_free_port(), True

    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
        try:
            sock.bind(("", preferred))
        except OSError:
            return _find_free_port(), True
        else:
            return preferred, False


@functools.lru_cache(maxsize=4)
def _load_model_cached(model_identifier: str) -> Any:
    resolved_identifier = model_identifier.strip() or DEFAULT_MODEL

    script_path = os.getenv(MODEL_SCRIPT_ENV)
    if script_path:
        custom_script = Path(script_path).expanduser()
        if not custom_script.exists():
            raise gr.Error(f"Custom modeling script not found: {custom_script}")

        module_name = f"open_provence_custom_{custom_script.stem}"
        spec = importlib.util.spec_from_file_location(module_name, custom_script)
        if spec is None or spec.loader is None:
            raise gr.Error(f"Unable to load module from {custom_script}")

        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        spec.loader.exec_module(module)  # type: ignore[call-arg]
        model_cls = getattr(module, "OpenProvenceModel", None)
        if model_cls is None:
            raise gr.Error("OpenProvenceModel class not found in custom script")

        model = model_cls.from_pretrained(resolved_identifier, trust_remote_code=True)
    else:
        model = AutoModel.from_pretrained(resolved_identifier, trust_remote_code=True)

    try:
        model = model.to(device=INFERENCE_DEVICE)  # type: ignore[attr-defined]
    except AttributeError:
        pass
    return model


def _format_summary(
    *,
    model_name: str,
    input_chars: int,
    output_chars: int,
    score: float | None,
    compression: float | None,
    duration: float | None,
) -> str:
    compression_percent = f"{compression:.1f}%" if isinstance(compression, (int, float)) else "N/A"
    score_str = f"{score:.4f}" if isinstance(score, (int, float)) else "N/A"
    duration_str = f"{duration:.2f}s" if isinstance(duration, (int, float)) else "N/A"
    return (
        f"**Model**: `{model_name}`\n"
        f"- Characters: {input_chars} → {output_chars}\n"
        f"- Score: {score_str}\n"
        f"- Compression: {compression_percent}\n"
        f"- Processing time: {duration_str}"
    )


def _flatten_first_sequence(value: Any) -> list[Any]:
    if value is None:
        return []
    if isinstance(value, (list, tuple)):
        if value and all(isinstance(item, (list, tuple)) for item in value):
            return list(value[0])
        return list(value)
    return [value]


def _render_sentence_scores_html(
    probabilities: list[float],
    kept_sentences: list[str],
    removed_sentences: list[str],
    threshold: float,
) -> str:
    if not probabilities:
        return (
            '<div class="sentence-score-wrapper empty">'
            "Sentence-level scores unavailable for this run."
            "</div>"
        )

    keep_flags = [float(prob) > threshold for prob in probabilities]
    keep_index = 0
    remove_index = 0
    rendered: list[tuple[str, float, bool]] = []

    for keep_flag, prob in zip(keep_flags, probabilities):
        if keep_flag:
            sentence = kept_sentences[keep_index] if keep_index < len(kept_sentences) else ""
            keep_index += 1
        else:
            sentence = (
                removed_sentences[remove_index] if remove_index < len(removed_sentences) else ""
            )
            remove_index += 1
        rendered.append((sentence, float(prob), keep_flag))

    style = """
<style>
.sentence-score-wrapper {
  display: flex;
  flex-direction: column;
  gap: 0.5rem;
}
.sentence-block {
  padding: 10px 12px;
  display: flex;
  gap: 10px;
  align-items: flex-start;
}
.sentence-block.keep {
  border-color: #a7f3d0;
}
.sentence-block.drop {
  border-color: #fecdd3;
}
.sentence-block .score-label {
  display: inline-block;
  padding: 4px 10px;
  border-radius: 999px;
  font-family: monospace;
  font-weight: 700;
  min-width: 72px;
  text-align: center;
  color: #166534;
  background: #dcfce7;
  border: 1px solid #a7f3d0;
}
.sentence-block.drop .score-label {
  color: #b91c1c;
  background: #ffe4e6;
  border-color: #fecdd3;
}
.sentence-block .text {
  white-space: pre-wrap;
  word-break: break-word;
  flex: 1;
}
</style>
""".strip()

    blocks: list[str] = []
    for sentence, prob, keep_flag in rendered:
        score_value = max(0.0, min(float(prob), 1.0))
        score_str = f"{score_value:.4f}"
        status_class = "keep" if keep_flag else "drop"
        safe_text = html.escape(sentence.strip() or "(empty sentence)")
        blocks.append(
            f'<div class="sentence-block {status_class}">'
            f'<span class="score-label">{score_str}</span>'
            f'<div class="text">{safe_text}</div>'
            "</div>"
        )

    return style + '<div class="sentence-score-wrapper">' + "".join(blocks) + "</div>"


def run_inference(
    model_identifier: str,
    threshold_value: float,
    question: str,
    title: str,
    text: str,
    *,
    progress: gr.Progress = gr.Progress(track_tqdm=False),
) -> tuple[str, str, str]:
    question_value = (question or "").strip()
    text_value = (text or "").strip()
    title_value = (title or "").strip() or None

    if not question_value:
        raise gr.Error("Please enter a question.")
    if not text_value:
        raise gr.Error("Please provide the context text.")

    model_name = (model_identifier or DEFAULT_MODEL).strip() or DEFAULT_MODEL

    progress(0.0, desc="Loading model…")

    try:
        model = _load_model_cached(model_name)
    except Exception as exc:  # pragma: no cover - user provided paths can fail
        raise gr.Error(f"Failed to load the model: {exc}") from exc

    clamped_threshold = min(max(float(threshold_value), 0.0), 1.0)

    process_kwargs: dict[str, Any] = {
        "question": question_value,
        "context": text_value,
        "threshold": clamped_threshold,
        "batch_size": DEFAULT_BATCH_SIZE,
        "language": "auto",
        "return_sentence_metrics": True,
        "return_sentence_texts": True,
    }
    if title_value:
        process_kwargs["title"] = title_value

    progress(0.5, desc="Running inference…")

    with torch.inference_mode():
        try:
            raw = model.process(**process_kwargs)  # type: ignore[attr-defined]
        except Exception as exc:  # pragma: no cover - external code
            raise gr.Error(f"An error occurred during inference: {exc}") from exc

    summary = _format_summary(
        model_name=model_name,
        input_chars=len(text_value),
        output_chars=len(_extract_string(raw.get("pruned_context"))),
        score=_extract_float(raw.get("reranking_score")),
        compression=_extract_float(raw.get("compression_rate")),
        duration=_extract_total_seconds(raw.get("timing"), raw.get("performance_trace")),
    )

    probabilities = [
        float(value)
        for value in _flatten_first_sequence(raw.get("sentence_probabilities"))
        if isinstance(value, (int, float))
    ]
    kept_sentences = [
        str(value) for value in _flatten_first_sequence(raw.get("kept_sentences")) if value
    ]
    removed_sentences = [
        str(value) for value in _flatten_first_sequence(raw.get("removed_sentences")) if value
    ]

    sentence_scores_html = _render_sentence_scores_html(
        probabilities=probabilities,
        kept_sentences=kept_sentences,
        removed_sentences=removed_sentences,
        threshold=clamped_threshold,
    )

    progress(1.0, desc="Completed")

    return summary, _extract_string(raw.get("pruned_context")), sentence_scores_html


def _extract_float(value: Any) -> float | None:
    if value is None:
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _extract_string(value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, str):
        return value
    if isinstance(value, (list, tuple)):
        parts = [_extract_string(v) for v in value if v]
        return "\n".join(part for part in parts if part)
    return str(value)


def _extract_total_seconds(timing: Any, performance_trace: Any) -> float | None:
    if isinstance(timing, dict):
        maybe = _extract_float(timing.get("total_seconds"))
        if maybe is not None:
            return maybe
    if isinstance(performance_trace, dict):
        return _extract_float(performance_trace.get("total_seconds"))
    return None


def build_interface() -> gr.Blocks:
    with gr.Blocks(title="OpenProvence Context Pruner & Reranker Demo") as demo:
        gr.Markdown(
            """
            ## ✂️ Open Provence Context Pruner & Reranker

            This demo showcases the
            [OpenProvence](https://github.com/hotchpotch/open_provence)
            project, which applies the [Provence](https://arxiv.org/abs/2501.16214)
            approach to remove question-irrelevant context
            while simultaneously computing reranker scores.
            Modern AI agents and context-engineering pipelines often recurse through
            large batches of search results.
            Irrelevant passages inflate the tokens we send to an LLM.
            By inserting a compact pre-processing model such as Provence,
            we pass along only the evidence that matters before the LLM generates an answer.

            OpenProvence is trained on public datasets with fully open training and inference code,
            and the weights are released under the MIT License.

            ### Demo Highlights 🚀
            - ⚡️ [30M-parameter compact model][hf-model]
              handles both English and Japanese on CPU -- lightning-fast on GPU.
            - ⚙️ Additional model families are available.
              Swap the model ID to try a different checkpoint.

            [hf-model]: https://huggingface.co/hotchpotch/open-provence-reranker-xsmall-v1
            """
        )

        with gr.Accordion("Help (click to expand)", open=False):
            gr.Markdown(
                """
                - **Model** 🔗: Provide a local path or Hugging Face model ID
                  that points to an OpenProvence checkpoint.
                - **threshold (recommend: 0.05-0.5)**: Range 0.00-1.00.
                  Higher values prune more aggressively. Lower values favor recall.
                  Adjust to suit your task.
                - **Title**: Supplying a concise title can improve pruning quality
                  when the context is long.
                - Sentence segmentation and language detection run in-model --
                  no manual tuning needed.
                - For implementation notes and evaluation results, see the
                  [Open Provence documentation](https://github.com/hotchpotch/open_provence).
                """
            )

        with gr.Row():
            model_input = gr.Textbox(
                label="Model",
                value=DEFAULT_MODEL,
                placeholder="e.g. hotchpotch/open-provence-reranker-xsmall-v1",
            )
            threshold_slider = gr.Slider(
                label="threshold (recommend: 0.05-0.5)",
                minimum=0.0,
                maximum=1.0,
                step=0.01,
                value=DEFAULT_THRESHOLD,
            )

        with gr.Group():
            question_input = gr.Textbox(
                label="Question (required)",
                lines=1,
            )
            title_input = gr.Textbox(
                label="Title (optional)",
                lines=1,
            )
            text_input = gr.Textbox(
                label="Context (required, long passages supported)",
                lines=12,
            )

        gr.Examples(
            examples=[[item.question, item.title, item.text] for item in EXAMPLE_INPUTS],
            inputs=[question_input, title_input, text_input],
            label="Sample input",
            examples_per_page=len(EXAMPLE_INPUTS),
        )

        run_button = gr.Button("Run inference", variant="primary")

        with gr.Tabs():
            with gr.Tab("Pruned Context"):
                summary_output = gr.Markdown(label="Summary")
                pruned_output = gr.Textbox(label="Pruned Text", lines=12)
            with gr.Tab("Sentence Scores"):
                sentence_scores_output = gr.HTML(
                    label="Sentence Scores",
                    value="Run inference to see per-sentence scores.",
                )

        run_button.click(
            fn=run_inference,
            inputs=[
                model_input,
                threshold_slider,
                question_input,
                title_input,
                text_input,
            ],
            outputs=[summary_output, pruned_output, sentence_scores_output],
        )

    return demo


demo = build_interface()


if __name__ == "__main__":
    requested_port = int(os.getenv("GRADIO_PORT", os.getenv("PORT", "7860")))
    port_value, port_was_adjusted = _resolve_server_port(requested_port)
    local_ips = _discover_local_ips()
    if local_ips:
        print("Local network URLs:")
        for ip in local_ips:
            print(f"  http://{ip}:{port_value}")
    else:
        print("No external IPv4 address detected; falling back to localhost only.")

    if port_was_adjusted and port_value != requested_port:
        print(f"Port {requested_port} unavailable; switched to {port_value}.")

    demo.launch(server_name="0.0.0.0", server_port=port_value, debug=True)