import os
import re
import time
from threading import Thread
from typing import Any, Dict, Generator, List, Tuple

import gradio as gr
import spaces
import torch
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

from nanogentzen.kernel import Sequent, verify_proof_tree
from nanogentzen.model import GentzenPolicyValueNet, PolicyValueConfig
from nanogentzen.parser import parse_natural_language, parse_symbolic_sequent
from nanogentzen.search import NeuralProofSearch
from nanogentzen.tokenizer import LogicTokenizer

# =====================================================================
# 1. LOAD LOCAL LLM (SYSTEM 1)
# =====================================================================
MODEL_ID = "Qwen/Qwen3.5-4B"

print(f"[*] Loading System 1 LLM: {MODEL_ID}...")
llm_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
llm_model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
llm_model.eval()
print("[+] System 1 loaded successfully.")

# =====================================================================
# 2. LOAD NANOGENTZEN KERNEL (SYSTEM 2)
# =====================================================================
def render_tree(node: dict, indent: int = 0) -> str:
    space = "  " * indent
    rule = node.get("rule", "UNKNOWN")
    seq = node.get("sequent", "")
    res = f"{space}• [{rule}] {seq}\n"
    for child in node.get("branches", []):
        res += render_tree(child, indent + 1)
    return res


def load_gentzen_engine():
    device = "cpu"
    tokenizer = LogicTokenizer()
    config = PolicyValueConfig(vocab_size=tokenizer.vocab_size)

    weights_file = (
        "nanogentzen_model.safetensors"
        if os.path.exists("nanogentzen_model.safetensors")
        else "model.safetensors"
    )
    if not os.path.exists(weights_file):
        raise FileNotFoundError(f"Missing weights file: {weights_file}")

    model = GentzenPolicyValueNet(config).to(device)
    model.load_state_dict(load_file(weights_file, device=device))
    model.eval()

    return NeuralProofSearch(model, tokenizer, device=device)


searcher = load_gentzen_engine()


def prove_comprehensive(query_str: str, max_depth: int = 8) -> Dict[str, Any]:
    seq = parse_symbolic_sequent(query_str)
    desc = "Formal symbolic sequent"
    if seq is None:
        nl_res = parse_natural_language(query_str)
        if nl_res:
            seq, desc = nl_res

    if seq is None:
        return {"success": False, "error": "Unable to parse into a valid sequent Γ ⊢ Δ."}

    # 1. Intuitionistic Logic (LI)
    t0 = time.perf_counter()
    li_tree = searcher.prove(seq, max_depth=max_depth)
    li_ms = round((time.perf_counter() - t0) * 1000.0, 2)
    li_sound = li_tree is not None and verify_proof_tree(li_tree)

    # 2. Classical Logic (LK via Glivenko Translation: Gamma |- ~~Delta)
    gamma_str = ", ".join(f.to_str() for f in seq.gamma) if seq.gamma else "0"
    delta_str = seq.delta[0].to_str() if seq.delta else "0"
    glivenko_seq = parse_symbolic_sequent(f"{gamma_str} |- ~~({delta_str})")

    t1 = time.perf_counter()
    lk_tree = searcher.prove(glivenko_seq, max_depth=max_depth) if glivenko_seq else None
    lk_ms = round((time.perf_counter() - t1) * 1000.0, 2)
    lk_sound = lk_tree is not None and verify_proof_tree(lk_tree)

    return {
        "success": True,
        "description": desc,
        "sequent": seq.to_str(),
        "li_proven": li_sound,
        "li_derivation": render_tree(li_tree) if li_sound else None,
        "li_latency_ms": li_ms,
        "lk_proven": lk_sound,
        "lk_derivation": render_tree(lk_tree) if lk_sound else None,
        "lk_latency_ms": lk_ms,
    }


def audit_neurosymbolic(think_text: str, user_input: str, max_depth: int = 8) -> Dict[str, Any]:
    has_logic = bool(re.search(r"(\|-|⟶|⊢|=>|&|~|\||\bif\b|\bthen\b|\btherefore\b)", user_input.lower()))
    proof_res = None

    if think_text:
        seq_match = re.search(
            r"(?:Formal Sequent:?\s*|Sequent:?\s*)?([A-Za-z0-9_~\(\)\s,=&|=>\-\>⟹∧∨¬]+(\|-|⟶|⊢)[A-Za-z0-9_~\(\)\s,=&|=>\-\>⟹∧∨¬]+)",
            think_text,
            re.IGNORECASE,
        )
        if seq_match:
            cand = seq_match.group(1).split("\n")[0].strip()
            alt = prove_comprehensive(cand, max_depth=max_depth)
            if alt.get("success"):
                proof_res = alt

    if (not proof_res or not proof_res.get("success")) and has_logic:
        nl_res = prove_comprehensive(user_input, max_depth=max_depth)
        if nl_res.get("success"):
            proof_res = nl_res

    if proof_res and proof_res.get("success"):
        proof_res["is_formal_proof"] = True
        return proof_res

    return {
        "success": True,
        "is_formal_proof": False,
        "li_proven": False,
        "lk_proven": False,
        "li_latency_ms": 1.5,
        "lk_latency_ms": 1.5,
    }

# =====================================================================
# 3. CHAT STREAMING PIPELINE (ZEROGPU)
# =====================================================================
SYSTEM_PROMPT = """You are an advanced Neurosymbolic AI Assistant.
Always start your response with a transparent, structured <think> block:
<think>
1. Problem Deconstruction: Identify premises and core goal.
2. Formal Sequent: If testing a deduction, write the symbolic sequent:
   Formal Sequent: (P => Q), P |- Q
3. Step-by-Step Analysis: Verify implications and rule applications.
</think>
After </think>, deliver a direct, comprehensive explanation."""


@spaces.GPU(duration=60)
def chat_stream(
    message: str,
    history: List[Dict[str, str]],
    temperature: float,
    max_tokens: int,
    max_depth: int,
) -> Generator[List[Dict[str, str]], None, None]:
    # 1. Direct formal sequent check
    has_turnstile = bool(re.search(r"(\|-|⟶|⊢)", message))
    if has_turnstile:
        res = prove_comprehensive(message, max_depth=int(max_depth))
        if res.get("success"):
            derivation = f"\n```text\n{res['li_derivation']}```" if res.get("li_derivation") else ""
            response = (
                f"### 🛡️ Sequent Certificate: `{res['sequent']}`\n\n"
                f"- **Intuitionistic Status (LI)**: {'✅ **PROVEN (Sound)**' if res['li_proven'] else '❌ **UNPROVABLE**'} (`{res['li_latency_ms']}ms`)\n"
                f"{derivation}\n"
                f"- **Classical Status (LK)**: {'✅ **VALID (Classical Tautology)**' if res['lk_proven'] else '❌ **INVALID**'} (`{res['lk_latency_ms']}ms`)"
            )
        else:
            response = f"❌ **Syntax / Parse Error**: {res.get('error')}"

        history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": response})
        yield history
        return

    # 2. Local Transformers Generation
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for msg in history:
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": message})

    prompt_text = llm_tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    model_inputs = llm_tokenizer([prompt_text], return_tensors="pt").to(llm_model.device)

    streamer = TextIteratorStreamer(
        llm_tokenizer,
        timeout=40.0,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    generate_kwargs = dict(
        model_inputs,
        streamer=streamer,
        max_new_tokens=int(max_tokens),
        temperature=float(temperature) if temperature > 0.0 else None,
        do_sample=temperature > 0.0,
        pad_token_id=llm_tokenizer.eos_token_id,
    )

    thread = Thread(target=llm_model.generate, kwargs=generate_kwargs)
    thread.start()

    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": ""})

    accumulated = ""
    for new_text in streamer:
        accumulated += new_text
        history[-1]["content"] = accumulated
        yield history

    # Extract thought trace and trigger System 2 verification
    think_body = ""
    if "<think>" in accumulated and "</think>" in accumulated:
        think_body = accumulated.split("</think>")[0].replace("<think>", "").strip()

    audit = audit_neurosymbolic(think_body, message, max_depth=int(max_depth))
    if audit.get("is_formal_proof") and audit.get("li_proven"):
        badge = f"\n\n---\n**⚡ nanoGentzen Audit (`{audit['li_latency_ms']}ms`):** ✅ `PROVEN SOUND (Q.E.D.)`\n"
        if audit.get("li_derivation"):
            badge += f"```text\n{audit['li_derivation']}```"
        accumulated += badge
    elif audit.get("is_formal_proof") and audit.get("lk_proven"):
        accumulated += f"\n\n---\n**⚡ nanoGentzen Audit (`{audit['lk_latency_ms']}ms`):** 🏛️ `CLASSICAL TAUTOLOGY (LK via Glivenko)`"
    elif audit.get("is_formal_proof") and not audit.get("li_proven") and not audit.get("lk_proven"):
        accumulated += f"\n\n---\n**⚡ nanoGentzen Audit (`{audit['li_latency_ms']}ms`):** ⚠️ `FALLACY DETECTED (Pruned Non-Sequitur)`"

    history[-1]["content"] = accumulated
    yield history


def prove_lab(sequent_str: str, max_depth: int) -> Tuple[str, str, str, str]:
    res = prove_comprehensive(sequent_str, max_depth=int(max_depth))
    if not res.get("success"):
        return "❌ Error", res.get("error", "Parse error"), "❌ Error", "Parse error"

    li_status = f"{'✅ PROVEN SOUND' if res['li_proven'] else '❌ UNPROVABLE'} ({res['li_latency_ms']}ms)"
    li_tree = res.get("li_derivation") or "No constructive proof tree."

    lk_status = f"{'✅ CLASSICAL TAUTOLOGY' if res['lk_proven'] else '❌ INVALID'} ({res['lk_latency_ms']}ms)"
    lk_tree = res.get("lk_derivation") or "Counter-model exists in Boolean valuation."

    return li_status, li_tree, lk_status, lk_tree

# =====================================================================
# 4. GRADIO LAYOUT
# =====================================================================
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 nanoGentzen Neurosymbolic Studio")
    gr.Markdown(
        f"**System 1 (`{MODEL_ID}` on ZeroGPU)** generates transparent `<think>` traces, while **System 2 (nanoGentzen)** certifies mathematical soundness in < 25 ms."
    )

    with gr.Tabs():
        with gr.Tab("💬 Neurosymbolic Chat"):
            chatbot = gr.Chatbot(height=520)
            with gr.Row():
                msg_input = gr.Textbox(
                    placeholder="Ask a logic question or enter a sequent (e.g., '(P => Q), ~Q |- ~P')...",
                    show_label=False,
                    scale=8,
                )
                submit_btn = gr.Button("Send 🚀", scale=1, variant="primary")

            with gr.Accordion("⚙️ Generation Hyperparameters", open=False):
                with gr.Row():
                    temp_slide = gr.Slider(0.0, 1.0, value=0.2, step=0.05, label="Temperature")
                    token_slide = gr.Slider(256, 2048, value=1024, step=128, label="Max Tokens")
                    depth_slide = gr.Slider(4, 16, value=8, step=1, label="nanoGentzen Max Depth")

            gr.Examples(
                examples=[
                    ["If it rains, the street is wet. The street is not wet. Did it rain?"],
                    ["((P => Q) => P) |- P"],
                    ["(P => Q), Q |- P"],
                    ["0 |- ~~ (~~P => P)"],
                ],
                inputs=msg_input,
            )

            submit_btn.click(
                chat_stream,
                inputs=[msg_input, chatbot, temp_slide, token_slide, depth_slide],
                outputs=[chatbot],
            ).then(lambda: "", None, msg_input)

            msg_input.submit(
                chat_stream,
                inputs=[msg_input, chatbot, temp_slide, token_slide, depth_slide],
                outputs=[chatbot],
            ).then(lambda: "", None, msg_input)

        with gr.Tab("🔬 Dual-Mode Logic Prover Lab"):
            gr.Markdown("### Interactive Sequent Verification ($LI$ vs $LK$)")
            with gr.Row():
                prover_input = gr.Textbox(
                    value="(P => Q), ~Q |- ~P",
                    label="Enter Sequent Γ ⊢ Δ",
                    scale=4,
                )
                prove_depth = gr.Slider(4, 16, value=8, step=1, label="Max Depth", scale=1)
                prove_btn = gr.Button("⚡ Prove Sequent", variant="primary", scale=1)

            with gr.Row():
                with gr.Column():
                    gr.Markdown("#### 🌿 Intuitionistic Logic (LI)")
                    li_status_out = gr.Textbox(label="Status", interactive=False)
                    li_tree_out = gr.Code(label="Derivation Tree", language="markdown")
                with gr.Column():
                    gr.Markdown("#### 🏛️ Classical Logic (LK via Glivenko)")
                    lk_status_out = gr.Textbox(label="Status", interactive=False)
                    lk_tree_out = gr.Code(label="Derivation Tree", language="markdown")

            prove_btn.click(
                prove_lab,
                inputs=[prover_input, prove_depth],
                outputs=[li_status_out, li_tree_out, lk_status_out, lk_tree_out],
            )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(primary_hue="indigo"))