"""Modelifier. Point at a model, add water, test the cut, create the model."""

import json
import os

import gradio as gr

from modelifier.surgeon import inspect_model, publish, test_edit

INK = "#0b0b0c"
PAPER = "#e7e2d6"
SIGNAL = "#7dd3c0"
MUTED = "#8b9098"
LINE = "#2c2c31"

CSS = f"""
body, .gradio-container {{ background: {INK}; color: {PAPER}; }}
.gradio-container {{ max-width: 980px; margin: 0 auto; }}
footer {{ display: none !important; }}
"""

THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.teal,
    secondary_hue=gr.themes.colors.stone,
    neutral_hue=gr.themes.colors.stone,
    font=gr.themes.GoogleFont("IBM Plex Sans"),
    font_mono=gr.themes.GoogleFont("IBM Plex Mono"),
).set(body_background_fill=INK, body_text_color=PAPER, block_background_fill="#141416")


def _layers(value: str):
    value = (value or "").strip()
    if not value:
        return None
    return [int(part.strip()) for part in value.split(",") if part.strip()]


def _water(value: str):
    return [line.strip() for line in (value or "").splitlines() if line.strip()]


def run_inspect(pointer: str, token: str):
    try:
        report = inspect_model(pointer, token.strip() or None)
    except Exception as exc:
        raise gr.Error(str(exc)) from exc
    lines = [
        f"{report['repo_id']}",
        f"files {report['files']} · hidden {report['hidden_size']} · layers {report['layers']}",
        f"embedding {report['embedding']}",
        "weights not downloaded",
        "",
    ]
    for row in report["resolved"]:
        lines.append(f"L{row['layer']}  {row['tensor']}  {row['shard']}")
    return "\n".join(lines), json.dumps(report, indent=2)


def run_test(pointer, water, note, layers, alpha, top_k, token):
    try:
        result = test_edit(
            pointer,
            _water(water),
            note or "",
            _layers(layers),
            float(alpha),
            int(top_k),
            token.strip() or None,
        )
    except Exception as exc:
        raise gr.Error(str(exc)) from exc
    lines = [
        f"{result['source']}  ·  not published",
        f"water samples {result['water_count']}  ·  alpha {result['alpha']}  ·  expected scale {1 - result['alpha']:.2f}",
        "",
    ]
    for row in result["layers"]:
        lines.append(
            f"L{row['layer']} {row['tensor']}  cols {row['columns']}  "
            f"scale {row['component_scale']}  ΔF {row['relative_frobenius_pct']}%  "
            f"σ2/σ1 {row['sigma2_over_sigma1_pct']}%"
        )
    lines.append("")
    lines.append("sha changed on every touched tensor. nothing was created.")
    return "\n".join(lines), json.dumps(result, indent=2)


def run_publish(pointer, water, note, layers, alpha, top_k, name, chat, private, token):
    if not token.strip():
        raise gr.Error("A write token is required to create the model. Test does not need one.")
    if not name.strip():
        raise gr.Error("Name the model.")
    try:
        result = publish(
            pointer,
            _water(water),
            name.strip(),
            note or "",
            _layers(layers),
            float(alpha),
            int(top_k),
            bool(chat),
            bool(private),
            token.strip(),
        )
    except Exception as exc:
        raise gr.Error(str(exc)) from exc
    lines = [
        result["model_url"],
        result.get("space_url") or "chat Space not requested",
        f"account {result['account']}",
    ]
    return "\n".join(lines), json.dumps({k: result[k] for k in ("model", "model_url", "space_url", "layers")}, indent=2)


with gr.Blocks(theme=THEME, css=CSS, title="Modelifier") as demo:
    gr.Markdown(
        """
# Modelifier
Point at a model. Point at datasets or other models as water. The cut is
`wⱼ ← wⱼ − α (uᵀ wⱼ) u` on the residual out-projection. Test writes nothing.
Create needs a write token and makes the named model, plus a chat Space if you ask.
        """
    )
    pointer = gr.Textbox(label="Model", placeholder="https://huggingface.co/org/model")
    water = gr.Textbox(
        label="Water — one model or dataset URL per line",
        lines=4,
        placeholder="https://huggingface.co/datasets/org/set",
    )
    note = gr.Textbox(label="What you want in the water", lines=2)
    with gr.Row():
        layers = gr.Textbox(label="Layers (blank = middle)", placeholder="15, 19")
        alpha = gr.Slider(0, 12, value=1, step=0.5, label="α  ·  1 removes, 6 flips")
        top_k = gr.Slider(4, 256, value=64, step=4, label="Columns K")
    token = gr.Textbox(label="Write token — create only", type="password")
    with gr.Row():
        inspect_btn = gr.Button("Read the model")
        test_btn = gr.Button("Test the cut", variant="primary")
    out = gr.Textbox(label="Result", lines=10)
    detail = gr.Code(label="Audit", language="json")
    inspect_btn.click(run_inspect, [pointer, token], [out, detail])
    test_btn.click(run_test, [pointer, water, note, layers, alpha, top_k, token], [out, detail])

    gr.Markdown("### Create under your account")
    with gr.Row():
        name = gr.Textbox(label="Model name", placeholder="qwen-terse")
        chat = gr.Checkbox(value=True, label="Also create a chat Space")
        private = gr.Checkbox(value=False, label="Private")
    create_btn = gr.Button("Create model", variant="primary")
    created = gr.Textbox(label="Created", lines=4)
    create_btn.click(
        run_publish,
        [pointer, water, note, layers, alpha, top_k, name, chat, private, token],
        [created, detail],
    )

if __name__ == "__main__":
    demo.launch()