import os
import json
import random
import gradio as gr

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# ======================================================
# ACE v6.0 — Lite-Core-Network (Phi-3 Mini local)
# - No evaluator agent (fixes podcast/meta nonsense)
# - Stronger meta filters
# - Stronger literal mode
# - Story-only narrative enforcement
# - ACW + CBS kept as "creative control" core
# ======================================================

MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct"
MEMORY_FILE = "ace_memory_v6_0.json"
MAX_NEW_TOKENS = 220  # we keep this (problems 2/10 intentionally NOT solved)

# -------- load model once (local, free) --------
device = 0 if torch.cuda.is_available() else -1

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
)
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=device,
)


# ======================================================
# Prompt classifiers
# ======================================================
def is_literal_mode(prompt: str) -> bool:
    p = prompt.lower()
    triggers = [
        "no metaphors",
        "without metaphors",
        "bez metafor",
        "bez prirovnaní",
        "literal explanation",
        "explain how",
        "explain in simple language",
        "explain simply",
        "explain this",
        "how does",
        "how do",
        "technical explanation",
        "mechanical clock",
        "mechanický",
        "mechanical",
    ]
    return any(t in p for t in triggers)


def is_short_greeting(prompt: str) -> bool:
    p = prompt.strip().lower()
    words = p.split()
    if len(words) <= 3 and any(
        w in p for w in ["hi", "hello", "hey", "yo", "sup", "čau", "čaute"]
    ):
        return True
    return False


def is_story_mode(prompt: str) -> bool:
    p = prompt.lower()
    triggers = [
        "story",
        "backstory",
        "character backstory",
        "short story",
        "scene",
        "narrative",
        "write a story",
        "sad but also has a good ending",
        "sad and ends hopeful",
        "tragic past",
        "origin story",
        "fantasy story",
        "sci-fi story",
        "horror story",
        "romantic story",
        "fairytale",
        "fairy tale",
    ]
    return any(t in p for t in triggers)


def wants_hopeful_ending(prompt: str) -> bool:
    p = prompt.lower()
    triggers = [
        "ends hopeful",
        "hopeful ending",
        "good ending",
        "ends with hope",
        "positive ending",
    ]
    return any(t in p for t in triggers)


def contains_figurative(text: str) -> bool:
    t = text.lower()
    figurative_triggers = [
        "like a",
        "as if",
        "as though",
        "seemed to",
        "resembled",
        "felt like",
        "heart of the",
        "soul of the",
        "spirit of the",
        "little people",
        "whispered to her soul",
        "whispered to his soul",
    ]
    return any(tr in t for tr in figurative_triggers)


def has_metadata_patterns(text: str) -> bool:
    t = text
    patterns = [
        "Title:",
        "Author:",
        "Publisher:",
        "[HOST]",
        "[GUEST]",
        "[AGENT]",
        "[USER]",
        "[ASSISTANT]",
        "ISBN",
        "Episode",
        "Podcast",
        "Top 10 tips",
    ]
    return any(pat in t for pat in patterns)


def has_meta_writing(text: str) -> bool:
    t = text.lower()
    patterns = [
        "to set the scene",
        "in this chapter",
        "in this story",
        "in your story",
        "as a writer",
        "you should",
        "the reader",
        "the audience",
        "this essay",
        "this chapter",
        "this section",
        "outline",
        "we will discuss",
        "top 10 tips",
        "in this episode",
        "in this podcast",
        "dear reader",
        "as a storyteller",
        "in the following story",
    ]
    return any(p in t for p in patterns)


def looks_truncated(text: str) -> bool:
    t = text.strip()
    if not t:
        return True
    last = t[-1]
    return last not in [".", "!", "?"]


def has_second_person_narration(text: str) -> bool:
    # catches "you / your" outside of quotes (rough heuristic)
    t = text.lower()
    return " you " in f" {t} " or " your " in f" {t} "


# ======================================================
# Context Boundary Space (CBS)
# ======================================================
def extract_context(prompt: str):
    words = [w.strip(".,!?").lower() for w in prompt.split()]
    keywords = {w for w in words if len(w) > 3}
    topic_hint = words[0] if words else ""
    sentiment = (
        "positive"
        if any(w in prompt.lower() for w in ["hope", "dream", "love", "happy"])
        else "negative"
        if any(w in prompt.lower() for w in ["fear", "pain", "loss", "sad", "hurt"])
        else "neutral"
    )
    return {
        "keywords": keywords,
        "topic_hint": topic_hint,
        "sentiment": sentiment,
    }


def context_fit(text: str, ctx) -> float:
    penalty = 0.0
    lowered = text.lower()

    if ctx["topic_hint"] and ctx["topic_hint"].lower() not in lowered:
        penalty += 0.15

    # simple keyword anchor: at least some overlap
    if ctx["keywords"]:
        overlap = sum(1 for k in ctx["keywords"] if k in lowered)
        if overlap == 0:
            penalty += 0.35
        elif overlap <= 2:
            penalty += 0.15

    if ctx["sentiment"] == "positive" and any(
        w in lowered for w in ["kill", "die", "ruin"]
    ):
        penalty += 0.35

    if ctx["sentiment"] == "negative" and any(
        w in lowered for w in ["cute", "joy", "happy", "celebrate"]
    ):
        penalty += 0.25

    return max(0.0, 1.0 - penalty)


# ======================================================
# ACW (Adaptive Creativity Window)
# ======================================================
def acw_state(
    prompt: str, history_len: int, decay: float, literal_mode: bool, story_mode: bool
) -> int:
    if literal_mode:
        return 0

    p = prompt.lower().strip()
    tokens = p.split()
    short = len(tokens) <= 4

    creative_trigger = story_mode or any(
        kw in p
        for kw in [
            "imagine",
            "poem",
            "invent",
            "creative",
            "world where",
            "dream",
            "weird idea",
            "make up",
        ]
    )

    if short and not creative_trigger:
        return 0

    entropy = random.uniform(0.45, 1.0)
    intensity = 0.85 if creative_trigger else 0.5
    stability = 1 - decay

    s = 0.45 * entropy + 0.35 * intensity + 0.2 * stability

    if s < 0.35:
        return 0
    elif s < 0.7:
        return 1
    else:
        return 2


def mutation_settings(state: int, literal_mode: bool, story_mode: bool):
    if literal_mode:
        return {"temp": 0.3, "top_p": 0.9, "count": 1}

    if story_mode:
        if state == 0:
            return {"temp": 0.9, "top_p": 0.95, "count": 2}
        if state == 1:
            return {"temp": 1.05, "top_p": 0.97, "count": 3}
        return {"temp": 1.15, "top_p": 0.98, "count": 3}

    if state == 0:
        return {"temp": 0.55, "top_p": 0.9, "count": 2}
    if state == 1:
        return {"temp": 0.95, "top_p": 0.94, "count": 3}
    return {"temp": 1.1, "top_p": 0.96, "count": 3}


# ======================================================
# Novelty + Twist scoring
# ======================================================
def rarity(text: str) -> float:
    words = text.lower().split()
    rare = [w for w in words if len(w) > 9]
    return min(len(rare) / 25.0, 1.0)


def metaphor_score(text: str) -> float:
    triggers = ["as if", "like a", "seemed to", "resembled", "as though"]
    return 1.0 if any(t in text.lower() for t in triggers) else 0.25


def twist(text: str, keywords) -> float:
    if not keywords:
        return 0.3
    lowered = text.lower()
    missing = sum(1 for k in keywords if k not in lowered)
    return min(missing / max(1, len(keywords)), 1.0)


def reward(text: str, ctx, literal_mode: bool, story_mode: bool) -> float:
    length_score = min(len(text) / 230.0, 1.0)
    novelty_score = rarity(text)
    imagery_score = metaphor_score(text)
    twist_score = twist(text, ctx["keywords"])
    context_score = context_fit(text, ctx)
    structure_score = max(0.0, 1.0 - 0.15 * text.count("\n\n"))

    creativity = 0.25 * novelty_score + 0.2 * imagery_score + 0.35 * twist_score

    base = (
        0.2 * length_score
        + 0.25 * creativity
        + 0.35 * context_score
        + 0.2 * structure_score
    )

    score = base

    if literal_mode:
        figurative_penalty = 0.9 if contains_figurative(text) else 0.0
        literal_score = (
            0.4 * length_score + 0.35 * context_score + 0.25 * structure_score
        )
        score = literal_score - figurative_penalty

    if story_mode:
        if has_meta_writing(text):
            score -= 1.2
        if has_metadata_patterns(text):
            score -= 1.0
    else:
        if has_metadata_patterns(text):
            score -= 0.8

    # bonus: punish second-person "you"/"your" in story mode (outside dialogue)
    if story_mode and has_second_person_narration(text):
        score -= 0.6

    return score


# ======================================================
# Memory
# ======================================================
def load_mem():
    if os.path.exists(MEMORY_FILE):
        with open(MEMORY_FILE, "r") as f:
            return json.load(f)
    return {"scores": [], "params": []}


def save_mem(mem):
    with open(MEMORY_FILE, "w") as f:
        json.dump(mem, f, indent=2)


# ======================================================
# LLM call — local Phi-3 via transformers
# ======================================================
def call_llm(prompt: str, max_new_tokens=MAX_NEW_TOKENS, temperature=0.7, top_p=0.9):
    try:
        out = generator(
            prompt,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
            pad_token_id=tokenizer.eos_token_id,
        )
        full_text = out[0]["generated_text"]
        completion = full_text[len(prompt):].strip()
        return completion
    except Exception as e:
        return f"[ACE ERROR] {e}"


# ======================================================
# Literal explainer
# ======================================================
def literal_explainer(prompt: str) -> str:
    full_prompt = (
        "Explain this as clearly, directly and literally as possible.\n"
        "No metaphors, no comparisons, no stories, no fictional books or podcasts.\n"
        "Use simple factual language only.\n\n"
        f"Question: {prompt}\n\nAnswer:"
    )
    return call_llm(full_prompt, max_new_tokens=220, temperature=0.3, top_p=0.9)


def sanitize_literal(text: str) -> str:
    if not contains_figurative(text):
        return text

    rewrite_prompt = (
        "Rewrite this explanation to be 100% literal.\n"
        "Remove ALL metaphors, comparisons, and figurative language.\n"
        "Use direct, simple, factual sentences only.\n\n"
        f"Original explanation:\n{text}\n\nLiteral rewrite:"
    )
    return call_llm(rewrite_prompt, max_new_tokens=200, temperature=0.25, top_p=0.9)


# ======================================================
# Story prompt builder + sanitizer
# ======================================================
def build_story_prompt(user_prompt: str, hopeful: bool) -> str:
    base = (
        "You are ACE, a narrative engine.\n"
        "Write a continuous story as pure narrative from inside the character's world.\n"
        "Do NOT explain writing techniques.\n"
        "Do NOT mention 'the reader', 'audience', 'chapter', 'section', or 'outline'.\n"
        "Do NOT give advice, steps, or tips.\n"
        "Just tell the story.\n"
    )
    if hopeful:
        base += (
            "The story should begin sad but end with a clear sense of hope, healing, or a new beginning.\n"
        )
    base += "\nStory prompt:\n" + user_prompt + "\n\nStory:\n"
    return base


def sanitize_story(text: str) -> str:
    # Hard filter for meta-writing / podcast artifacts
    if not (has_meta_writing(text) or has_metadata_patterns(text)):
        return text

    # Try to chop off leading meta paragraphs and keep the narrative part
    paragraphs = [p.strip() for p in text.split("\n") if p.strip()]
    kept = []
    for p in paragraphs:
        if has_meta_writing(p) or has_metadata_patterns(p):
            continue
        kept.append(p)
    if kept:
        return "\n\n".join(kept).strip()
    return text


# ======================================================
# Explorer (single agent in v6 Lite-Core)
# ======================================================
def explorer(prompt: str, temp: float, top_p: float) -> str:
    return call_llm(prompt, max_new_tokens=MAX_NEW_TOKENS, temperature=temp, top_p=top_p)


# ======================================================
# ACE v6.0 core
# ======================================================
def ace_generate(prompt: str, history, mode: str) -> str:
    ctx = extract_context(prompt)
    mem = load_mem()

    literal = is_literal_mode(prompt)
    short_greet = is_short_greeting(prompt)
    story_mode = is_story_mode(prompt)
    hopeful = wants_hopeful_ending(prompt)

    if mode == "Creative":
        literal = False

    if literal and mode != "Creative":
        base = literal_explainer(prompt)
        return sanitize_literal(base)

    if short_greet and mode != "Creative":
        return "Ahoj, som ACE v6.0 Lite-Core-Network. Čo chceš skúsiť?"

    decay = sum(mem["scores"][-5:]) / 5 if mem["scores"] else 0.0
    state = acw_state(prompt, len(history), decay, literal, story_mode)

    if mode == "Instant":
        temp = 0.4 if literal else (1.0 if story_mode else 0.9)
        top_p = 0.95 if story_mode else 0.9
        if story_mode:
            story_prompt = build_story_prompt(prompt, hopeful)
            raw = explorer(story_prompt, temp, top_p)
            cleaned = sanitize_story(raw)
            return cleaned
        else:
            return explorer(prompt, temp, top_p)

    if mode == "Full":
        settings = mutation_settings(state, literal, story_mode)
        use_memory = True

    elif mode == "Lite":
        settings = mutation_settings(state, literal, story_mode)
        settings["count"] = 1
        use_memory = True

    elif mode == "Turbo":
        settings = mutation_settings(state, literal, story_mode)
        settings["count"] = 2
        use_memory = True

    elif mode == "Creative":
        state = 2
        settings = mutation_settings(state, False, True if story_mode else False)
        settings["count"] = max(3, settings["count"])
        literal = False
        use_memory = True

    else:
        settings = mutation_settings(state, literal, story_mode)
        use_memory = True

    raw_candidates = []
    for _ in range(settings["count"]):
        if story_mode:
            story_prompt = build_story_prompt(prompt, hopeful)
            text = explorer(story_prompt, settings["temp"], settings["top_p"])
        else:
            text = explorer(prompt, settings["temp"], settings["top_p"])
        raw_candidates.append(text)

    improved_candidates = []
    for text in raw_candidates:
        if story_mode:
            text = sanitize_story(text)
        if literal:
            text = sanitize_literal(text)

        score = reward(text, ctx, literal, story_mode)
        improved_candidates.append((text, score))

    if not improved_candidates:
        fallback_temp = 0.35 if literal else (1.0 if story_mode else 0.9)
        fallback_top_p = 0.95 if story_mode else 0.9
        if story_mode:
            story_prompt = build_story_prompt(prompt, hopeful)
            best = explorer(story_prompt, fallback_temp, fallback_top_p)
            best = sanitize_story(best)
        else:
            best = explorer(prompt, fallback_temp, fallback_top_p)
        best_score = reward(best, ctx, literal, story_mode)
    else:
        best, best_score = max(improved_candidates, key=lambda x: x[1])

    if use_memory:
        mem["scores"].append(best_score)
        mem["params"].append(settings)
        mem["scores"] = mem["scores"][-80:]
        mem["params"] = mem["params"][-80:]
        save_mem(mem)

    return best


# ======================================================
# Gradio UI
# ======================================================
def ui_response(msg, history, mode):
    return ace_generate(msg, history, mode)


mode_dropdown = gr.Dropdown(
    choices=["Full", "Lite", "Turbo", "Instant", "Creative"],
    value="Full",
    label="ACE Mode",
)

demo = gr.ChatInterface(
    ui_response,
    additional_inputs=[mode_dropdown],
    title="ACE v6.0 — Lite-Core-Network (Phi-3 Mini, local)",
    description=(
        "ACCF-based creative engine with CBS, ACW, Lite-Core candidate network, "
        "story-only narrative mode, literal mode, and meta-writing suppression. "
        "Running locally on microsoft/Phi-3-mini-4k-instruct (CPU)."
    ),
)

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