from __future__ import annotations

import os
from pathlib import Path
from typing import Callable


LABELS = {
    0: "Normal",
    1: "Hate/Offensive",
}

WEIGHT_FILENAMES = ("model.safetensors", "pytorch_model.bin")
TOKENIZER_FILENAMES = (
    "tokenizer.json",
    "tokenizer_config.json",
    "vocab.txt",
    "special_tokens_map.json",
)

DEFAULT_MODEL_PATH = "/Users/qqq/Downloads/fine_tuned_hatebert_model"

LIGHT_KEYNOTE_CSS = """
:root {
  --paper: #f7fbff;
  --paper-2: #eef5fb;
  --paper-3: #e8eef5;
  --ink: #15202a;
  --muted: #5f6f7d;
  --accent: #4f87b3;
  --accent-soft: #d8e7f4;
  --clear: #2d7a57;
  --danger: #b65745;
  --line: rgba(35, 71, 102, 0.12);
}

.gradio-container {
  background:
    radial-gradient(circle at 86% 14%, rgba(107, 177, 226, 0.18), transparent 18%),
    linear-gradient(90deg, rgba(35, 71, 102, 0.05) 1px, transparent 1px),
    linear-gradient(rgba(35, 71, 102, 0.04) 1px, transparent 1px),
    linear-gradient(150deg, var(--paper) 0%, var(--paper-2) 46%, var(--paper-3) 100%) !important;
  background-size: auto, 40px 40px, 40px 40px, auto !important;
  color: var(--ink) !important;
  font-family: "Aptos", "Segoe UI", sans-serif !important;
}

.prototype-shell {
  max-width: 1320px;
  margin: 0 auto;
  padding: 18px 18px 28px;
}

.topbar {
  display: flex;
  justify-content: space-between;
  align-items: center;
  gap: 20px;
  margin-bottom: 16px;
}

.brand {
  display: flex;
  align-items: center;
  gap: 12px;
}

.brand-mark {
  width: 36px;
  height: 36px;
  border-radius: 12px;
  background: linear-gradient(135deg, #14212b, #36556e);
  color: #9fd0f7;
  display: grid;
  place-items: center;
  font-weight: 700;
  font-size: 15px;
}

.brand-copy {
  min-width: 0;
}

.brand-title {
  font-family: Georgia, "Times New Roman", serif;
  font-size: 19px;
}

.brand-subtitle {
  margin-top: 2px;
  color: rgba(21, 32, 42, 0.56);
  font-size: 11px;
  letter-spacing: 0.14em;
  text-transform: uppercase;
}

.status-badge {
  padding: 10px 12px;
  border: 1px solid rgba(35, 71, 102, 0.1);
  background: rgba(255, 255, 255, 0.52);
  color: #2e5b82;
  font-size: 11px;
  letter-spacing: 0.14em;
  text-transform: uppercase;
}

.prototype-stage {
  display: grid;
  grid-template-columns: minmax(0, 1.12fr) minmax(410px, 0.88fr);
  gap: 24px;
  align-items: start;
}

.stage-copy {
  padding: 10px 0 6px;
}

.stage-kicker {
  color: var(--accent);
  font-size: 12px;
  font-weight: 700;
  letter-spacing: 0.18em;
  text-transform: uppercase;
}

.stage-title {
  margin: 14px 0 0;
  max-width: 9em;
  font-family: Georgia, "Times New Roman", serif;
  font-size: clamp(42px, 5.2vw, 68px);
  line-height: 0.96;
  letter-spacing: -0.05em;
}

.stage-subtitle {
  max-width: 35rem;
  margin-top: 14px;
  color: var(--muted);
  font-size: 15px;
  line-height: 1.55;
}

.stage-proof {
  display: grid;
  grid-template-columns: repeat(3, minmax(0, 1fr));
  gap: 12px;
  margin-top: 18px;
  max-width: 700px;
}

.stage-proof-item {
  border-top: 1px solid var(--line);
  padding-top: 12px;
}

.stage-proof-label {
  color: var(--muted);
  font-size: 10px;
  letter-spacing: 0.14em;
  text-transform: uppercase;
}

.stage-proof-value {
  margin-top: 6px;
  font-family: Georgia, "Times New Roman", serif;
  font-size: 22px;
}

.stage-chips {
  margin-top: 14px;
  display: flex;
  gap: 8px;
  flex-wrap: wrap;
}

.stage-chip {
  padding: 8px 10px;
  border-radius: 999px;
  border: 1px solid rgba(37, 85, 123, 0.1);
  background: rgba(37, 85, 123, 0.06);
  color: #25557b;
  font-size: 10px;
  letter-spacing: 0.12em;
  text-transform: uppercase;
}

.queue {
  margin-top: 16px;
  display: grid;
  gap: 12px;
  max-width: 640px;
}

.queue-card {
  display: grid;
  grid-template-columns: 68px 1fr auto;
  gap: 14px;
  align-items: center;
  padding: 14px;
  border: 1px solid rgba(35, 71, 102, 0.08);
  background: rgba(255, 255, 255, 0.62);
  box-shadow: 0 16px 34px rgba(114, 145, 175, 0.08);
}

.queue-score {
  font-family: Georgia, "Times New Roman", serif;
  font-size: 24px;
  color: #25557b;
}

.queue-title {
  font-size: 12px;
  letter-spacing: 0.14em;
  text-transform: uppercase;
  color: rgba(21, 32, 42, 0.52);
}

.queue-text {
  margin-top: 4px;
  color: rgba(21, 32, 42, 0.78);
  line-height: 1.42;
}

.queue-pill {
  padding: 8px 10px;
  border-radius: 999px;
  border: 1px solid rgba(37, 85, 123, 0.1);
  background: rgba(37, 85, 123, 0.06);
  color: #25557b;
  font-size: 10px;
  letter-spacing: 0.12em;
  text-transform: uppercase;
}

.stage-note {
  margin-top: 14px;
  padding: 12px 14px;
  border-left: 3px solid #5a8db5;
  background: rgba(255, 255, 255, 0.48);
  color: rgba(21, 32, 42, 0.76);
  line-height: 1.5;
}

.prototype-panel {
  background: rgba(255, 255, 255, 0.58);
  border: 1px solid rgba(35, 71, 102, 0.1);
  box-shadow: 0 22px 50px rgba(96, 132, 166, 0.12);
  backdrop-filter: blur(16px);
  padding: 24px;
}

.panel-head {
  display: flex;
  justify-content: space-between;
  align-items: flex-start;
  gap: 14px;
}

.panel-kicker {
  color: rgba(21, 32, 42, 0.55);
  font-size: 11px;
  letter-spacing: 0.16em;
  text-transform: uppercase;
}

.panel-title {
  margin-top: 8px;
  font-family: Georgia, "Times New Roman", serif;
  font-size: 34px;
  line-height: 0.98;
}

.panel-badge {
  padding: 8px 10px;
  border: 1px solid rgba(79, 135, 179, 0.16);
  background: rgba(79, 135, 179, 0.08);
  color: #4076a1;
  font-size: 10px;
  letter-spacing: 0.14em;
  text-transform: uppercase;
}

.prototype-panel textarea,
.prototype-panel input,
.prototype-panel .gradio-textbox textarea {
  background: rgba(243, 249, 255, 0.82) !important;
  border: 1px solid rgba(35, 71, 102, 0.1) !important;
  color: rgba(21, 32, 42, 0.78) !important;
}

.prototype-panel label,
.prototype-panel .gradio-form label {
  color: rgba(21, 32, 42, 0.55) !important;
  font-size: 11px !important;
  letter-spacing: 0.14em !important;
  text-transform: uppercase !important;
}

.prototype-panel button.primary,
.prototype-panel button.lg.primary {
  background: linear-gradient(135deg, #1f4e72, #5a8db5) !important;
  border: 0 !important;
  color: white !important;
}

.prototype-panel button.secondary {
  border: 1px solid rgba(35, 71, 102, 0.1) !important;
  background: rgba(255, 255, 255, 0.54) !important;
  color: #234966 !important;
}

.prototype-result {
  margin-top: 22px;
  padding-top: 22px;
  border-top: 1px solid rgba(35, 71, 102, 0.1);
  display: grid;
  grid-template-columns: 1fr 168px;
  gap: 16px;
  align-items: center;
}

.result-copy {
  min-width: 0;
}

.result-kicker {
  color: rgba(21, 32, 42, 0.52);
  font-size: 10px;
  letter-spacing: 0.16em;
  text-transform: uppercase;
}

.result-label {
  margin-top: 8px;
  font-family: Georgia, "Times New Roman", serif;
  font-size: 48px;
  line-height: 0.95;
}

.result-confidence {
  margin-top: 10px;
  max-width: 17rem;
  color: rgba(21, 32, 42, 0.66);
  font-size: 15px;
  line-height: 1.55;
}

.result-ring {
  width: 160px;
  height: 160px;
  border-radius: 50%;
  display: grid;
  place-items: center;
  justify-self: end;
  background: conic-gradient(from -90deg, #5c93c0 0 0%, rgba(35, 71, 102, 0.08) 0% 100%);
}

.result-ring-inner {
  width: 114px;
  height: 114px;
  border-radius: 50%;
  background: rgba(246, 250, 255, 0.96);
  display: grid;
  place-items: center;
  text-align: center;
}

.result-ring-value {
  font-family: Georgia, "Times New Roman", serif;
  font-size: 30px;
}

.result-ring-label {
  margin-top: 4px;
  color: rgba(21, 32, 42, 0.54);
  font-size: 10px;
  letter-spacing: 0.14em;
  text-transform: uppercase;
}

.normal-result {
  --result-accent: var(--clear);
}

.harmful-result {
  --result-accent: var(--danger);
}

.error-result {
  --result-accent: var(--danger);
}

.neutral-result {
  --result-accent: var(--accent);
}

@media (max-width: 1100px) {
  .topbar {
    align-items: flex-start;
    flex-direction: column;
  }

  .prototype-stage {
    grid-template-columns: 1fr;
  }

  .queue {
    max-width: 100%;
  }
}

@media (max-width: 680px) {
  .prototype-result {
    grid-template-columns: 1fr;
  }

  .result-ring {
    justify-self: start;
  }
}
"""

Classifier = Callable[[str], tuple[str, float]]


def detect_missing_model_artifacts(model_dir: Path | str) -> list[str]:
    path = Path(model_dir)
    missing: list[str] = []

    if not (path / "config.json").exists():
        missing.append("config.json")

    if not any((path / name).exists() for name in WEIGHT_FILENAMES):
        missing.append("model weights (model.safetensors or pytorch_model.bin)")

    if not any((path / name).exists() for name in TOKENIZER_FILENAMES):
        missing.append("tokenizer assets")

    return missing


def build_colab_mount_hint() -> str:
    return (
        "Mount Google Drive first with: "
        "from google.colab import drive; drive.mount('/content/drive')"
    )


def build_model_path_hint(model_path: Path | str) -> str:
    path = str(model_path)
    if path.startswith("/content/drive/"):
        return build_colab_mount_hint()
    return (
        "Check that the local model path is correct or set "
        "HATEBERT_MODEL_PATH to your model directory."
    )


def get_bundled_model_path(base_dir: Path | str | None = None) -> Path:
    root = Path(base_dir) if base_dir is not None else Path(__file__).resolve().parent
    return root / "fine_tuned_hatebert_model"


def resolve_model_path(base_dir: Path | str | None = None) -> str:
    env_path = os.getenv("HATEBERT_MODEL_PATH")
    if env_path:
        return env_path

    bundled = get_bundled_model_path(base_dir)
    if bundled.exists():
        return str(bundled)

    return DEFAULT_MODEL_PATH


def get_transformers_load_kwargs() -> dict[str, bool]:
    return {"local_files_only": True}


def build_launch_kwargs() -> dict[str, object]:
    env_port = os.getenv("GRADIO_SERVER_PORT") or os.getenv("PORT")
    server_port = int(env_port) if env_port else None
    running_on_space = bool(os.getenv("SPACE_ID") or os.getenv("SPACE_HOST"))
    return {
        "share": not running_on_space,
        "debug": not running_on_space,
        "server_name": "0.0.0.0",
        "server_port": server_port,
        "show_error": True,
    }


def validate_model_path(model_path: Path | str) -> tuple[bool, str]:
    path = Path(model_path)

    if not path.exists():
        return False, f"Model path does not exist: {path}. {build_model_path_hint(path)}"

    if not path.is_dir():
        return False, f"Model path is not a directory: {path}"

    missing = detect_missing_model_artifacts(path)
    if missing:
        joined = ", ".join(missing)
        return False, f"Missing required files in model directory: {joined}"

    return True, f"Model directory ready: {path}"


def render_empty_state() -> str:
    return """
    <div class="prototype-result neutral-result">
      <div class="result-copy">
        <div class="result-kicker">Prediction</div>
        <div class="result-label">Awaiting input</div>
        <div class="result-confidence">Enter a short online comment to run inference.</div>
      </div>
      <div class="result-ring" style="background: conic-gradient(from -90deg, var(--accent) 0 0%, rgba(35, 71, 102, 0.08) 0% 100%);">
        <div class="result-ring-inner">
          <div>
            <div class="result-ring-value">0%</div>
            <div class="result-ring-label">Confidence</div>
          </div>
        </div>
      </div>
    </div>
    """


def render_error_state(message: str) -> str:
    return f"""
    <div class="prototype-result error-result">
      <div class="result-copy">
        <div class="result-kicker">Prediction unavailable</div>
        <div class="result-label">Startup issue</div>
        <div class="result-confidence">{message}</div>
      </div>
      <div class="result-ring" style="background: conic-gradient(from -90deg, var(--danger) 0 0%, rgba(35, 71, 102, 0.08) 0% 100%);">
        <div class="result-ring-inner">
          <div>
            <div class="result-ring-value">0%</div>
            <div class="result-ring-label">Confidence</div>
          </div>
        </div>
      </div>
    </div>
    """


def render_result_state(label: str, confidence: float) -> str:
    state_class = "normal-result" if label == "Normal" else "harmful-result"
    accent = "var(--clear)" if label == "Normal" else "var(--danger)"
    confidence_pct = confidence * 100
    return f"""
    <div class="prototype-result {state_class}">
      <div class="result-copy">
        <div class="result-kicker">Prediction</div>
        <div class="result-label">{label}</div>
        <div class="result-confidence">Model confidence: {confidence:.2%}</div>
      </div>
      <div class="result-ring" style="background: conic-gradient(from -90deg, {accent} 0 {confidence_pct:.2f}%, rgba(35, 71, 102, 0.08) {confidence_pct:.2f}% 100%);">
        <div class="result-ring-inner">
          <div>
            <div class="result-ring-value">{confidence:.2%}</div>
            <div class="result-ring-label">Confidence</div>
          </div>
        </div>
      </div>
    </div>
    """


def predict_text(text: str | None, classify: Classifier) -> str:
    cleaned = (text or "").strip()
    if not cleaned:
        return render_empty_state()

    try:
        label, confidence = classify(cleaned)
    except Exception as exc:
        return render_error_state(f"Inference unavailable: {exc}")

    return render_result_state(label, confidence)


def build_hero_markup() -> str:
    return """
    <div class="topbar">
      <div class="brand">
        <div class="brand-mark">H</div>
        <div class="brand-copy">
          <div class="brand-title">HateBERT Moderation Prototype</div>
          <div class="brand-subtitle">Platform Safety / Live Inference System</div>
        </div>
      </div>
      <div class="status-badge">Binary classifier / Real-time review</div>
    </div>

    <section class="prototype-stage">
      <div class="stage-copy">
        <div class="stage-kicker">Content safety intelligence</div>
        <h1 class="stage-title">Moderation infrastructure for live social risk screening.</h1>
        <p class="stage-subtitle">
          A polished live review surface for harmful content detection. This interface frames
          the model as a platform moderation tool rather than a plain classroom demo.
        </p>
        <div class="stage-proof">
          <div class="stage-proof-item">
            <div class="stage-proof-label">System</div>
            <div class="stage-proof-value">Moderation</div>
          </div>
          <div class="stage-proof-item">
            <div class="stage-proof-label">Model</div>
            <div class="stage-proof-value">HateBERT</div>
          </div>
          <div class="stage-proof-item">
            <div class="stage-proof-label">Mode</div>
            <div class="stage-proof-value">Live Review</div>
          </div>
        </div>

        <div class="stage-chips">
          <div class="stage-chip">Binary classifier</div>
          <div class="stage-chip">Confidence output</div>
          <div class="stage-chip">2 classes</div>
        </div>

        <div class="queue">
          <div class="queue-card">
            <div class="queue-score">82%</div>
            <div>
              <div class="queue-title">Flagged Comment</div>
              <div class="queue-text">User-submitted social content is routed into a lightweight review layer with confidence evidence.</div>
            </div>
            <div class="queue-pill">Queued</div>
          </div>

          <div class="queue-card">
            <div class="queue-score">2</div>
            <div>
              <div class="queue-title">Output Classes</div>
              <div class="queue-text">The interface returns a binary moderation decision designed for clear academic demonstration and platform governance storytelling.</div>
            </div>
            <div class="queue-pill">Normal / Harmful</div>
          </div>
        </div>

        <div class="stage-note">
          Designed to look and behave like a real platform safety product while staying lightweight enough for sharing and live presentation.
        </div>
      </div>
    </section>
    """


def load_runtime(model_path: str) -> Classifier:
    ok, message = validate_model_path(model_path)
    if not ok:
        raise RuntimeError(message)

    import torch
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    load_kwargs = get_transformers_load_kwargs()
    tokenizer = AutoTokenizer.from_pretrained(model_path, **load_kwargs)
    model = AutoModelForSequenceClassification.from_pretrained(model_path, **load_kwargs)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()

    def classify(cleaned_text: str) -> tuple[str, float]:
        inputs = tokenizer(
            cleaned_text,
            return_tensors="pt",
            truncation=True,
            padding=True,
            max_length=128,
        )
        inputs = {key: value.to(device) for key, value in inputs.items()}

        with torch.no_grad():
            logits = model(**inputs).logits
            probabilities = torch.softmax(logits, dim=-1)[0]
            predicted_id = int(torch.argmax(probabilities).item())
            confidence = float(probabilities[predicted_id].item())

        label = LABELS.get(predicted_id, f"LABEL_{predicted_id}")
        return label, confidence

    return classify


def build_demo(classify: Classifier):
    import gradio as gr

    with gr.Blocks(title="Moderation Demo", css=LIGHT_KEYNOTE_CSS) as demo:
        with gr.Column(elem_classes=["prototype-shell"]):
            gr.HTML(build_hero_markup())

            with gr.Column(elem_classes=["prototype-panel"]):
                gr.HTML(
                    """
                    <div class="panel-head">
                      <div>
                        <div class="panel-kicker">Inference Workspace</div>
                        <div class="panel-title">Live Comment Review</div>
                      </div>
                      <div class="panel-badge">Fine-tuned model</div>
                    </div>
                    """
                )

                text_input = gr.Textbox(
                    label="Comment Input",
                    lines=7,
                    placeholder="Enter an online comment for classification...",
                )
                with gr.Row():
                    analyze_button = gr.Button("Run Moderation", variant="primary")
                    clear_button = gr.Button("Clear", variant="secondary")
                output = gr.HTML(value=render_empty_state())

                gr.Examples(
                    examples=[
                        ["I really enjoy spending time with my friends at the park."],
                        ["You are absolutely useless and I hate everything about you."],
                        ["This is a neutral statement about the weather."],
                    ],
                    inputs=text_input,
                    label="Live Review Examples",
                )

                analyze_button.click(
                    fn=lambda text: predict_text(text, classify),
                    inputs=text_input,
                    outputs=output,
                )
                clear_button.click(
                    fn=lambda: ("", render_empty_state()),
                    outputs=[text_input, output],
                )
                text_input.submit(
                    fn=lambda text: predict_text(text, classify),
                    inputs=text_input,
                    outputs=output,
                )

    return demo


def main() -> None:
    """
    Local quick start:
    1. Put the fine-tuned model directory on this machine.
    2. Set HATEBERT_MODEL_PATH if your model folder is not the default path.
    3. Run this script to launch the Gradio demo.
    4. Share the temporary public link with your teammates.
    """

    classify = load_runtime(resolve_model_path())
    demo = build_demo(classify)
    demo.launch(**build_launch_kwargs())


if __name__ == "__main__":
    main()