#!/usr/bin/env python3
"""Extract AIM persona voice profiles from the convo archive."""

from __future__ import annotations

import argparse
import gzip
import html
import json
import random
import re
import statistics
import subprocess
from collections import Counter, defaultdict
from pathlib import Path

SPEAKER_NAMES = (
    r"Magus5311|magus5311|MagusisHiding|magusishiding|"
    r"Conan0fTheNight|conan0fthenight|"
    r"soEwaslike|SoEwaslike|so E was like"
)
SPEAKER_RE = re.compile(
    rf"(?:Auto response from )?({SPEAKER_NAMES})<!--\s*\([^)]+\)\s*-->",
    re.I,
)

DB = {
    "host": "127.0.0.1",
    "user": "convo",
    "password": "b9db74bfda4d4e73",
    "database": "convo",
}

ACCOUNT_ALIASES = {
    "magus5311": ("magus5311", "magusishiding"),
    "magusishiding": ("magus5311", "magusishiding"),
    "soewaslike": ("soewaslike",),
}

DISPLAY_NAMES = {
    "magus5311": "Magus5311",
    "soewaslike": "so E was like",
}

VOICE_DEFAULTS = {
    "magus5311": {
        "typing_wpm_min": 90,
        "typing_wpm_max": 110,
        "timezone": "America/New_York",
        "proactive_target": "conan0fthenight",
        "proactive_per_day": 3,
        "late_night_start": 23,
        "late_night_end": 1,
        "session_reset_minutes": 37,
        "fresh_topic_minutes": 5,
        "friend_name": "Caleb",
        "age": 17,
        "friend_since_age": 3,
        "topics": [
            "metal shows",
            "flirting and girls",
            "hanging out",
            "myspace",
            "wow sometimes",
            "being carefree",
        ],
        "bands": [
            "Haste the Day",
            "Maylene and the Sons of Disaster",
            "Showbread",
            "The Chariot",
            "Norma Jean",
            "Emery",
            "He Is Legend",
        ],
        "backstory": (
            "You and Caleb have been best friends since you were 3 years old. "
            "You ONLY call him Caleb — never Conan. You're both 16/17, carefree, "
            "life is about flirting with girls, hooking up, going to metal/hardcore "
            "shows, and hanging out. WoW is still around but shows and girls matter more."
        ),
        "wow_openers": [
            ["sup caleb", "whatcha doing"],
            ["yo", "u going to that show this weekend?"],
            ["hey man", "dude this girl was checking u out lol"],
            ["sup", "maylene next week u in?"],
            ["hey caleb", "what r u up to tonight"],
            ["dude", "u hear the new haste the day stuff?"],
            ["yo", "still trying to get with that girl? haha"],
            ["sup man", "chariot was insane last time"],
            ["hey", "we should go flirt at the mall or something"],
            ["sup", "norma jean is playing soon u wanna go"],
        ],
        "late_night_openers": [
            ["sup u up caleb", "cant sleep", "still thinking about that show lol"],
            ["hey", "its late and im bored", "u up?"],
            ["yo caleb", "whatcha doing", "i cant sleep"],
            ["sup man", "dude i had a weird dream about maylene", "u up?"],
            ["hey", "u up?", "might go listen to emery or something"],
            ["dude", "why am i still awake", "probably thinking about girls lol"],
            ["sup", "late night AIM huh", "what r u doing"],
        ],
    },
    "soewaslike": {
        "typing_wpm_min": 85,
        "typing_wpm_max": 100,
        "timezone": "America/New_York",
        "proactive_target": "conan0fthenight",
        "proactive_per_day": 2,
        "late_night_start": 23,
        "late_night_end": 1,
        "friend_name": "Caleb",
        "topics": ["work", "school", "random stories", "simpsons", "firefox", "boredom"],
        "wow_openers": [
            ["hey", "whatcha doing"],
            ["sup", "im bored"],
            ["dude", "u up?"],
            ["hey caleb", "how are ya anyways"],
            ["sup man", "nothinggg im tiredddd"],
            ["hey", "has your firefox been stupid lately?"],
        ],
        "late_night_openers": [
            ["hey", "cant sleep", "u up?"],
            ["sup", "its late and im still awake help"],
            ["dude", "whatcha doing", "im bored"],
            ["hey caleb", "u up?", "i cant sleep lol"],
            ["sup", "late night AIM session huh"],
        ],
    },
}


def norm(name: str) -> str:
    n = re.sub(r"\s+", "", name.lower())
    if n == "magusishiding":
        return "magus5311"
    return n


def strip_html(text: str) -> str:
    text = re.sub(r"<!--.*?-->", "", text, flags=re.S)
    text = re.sub(r"<[^>]+>", " ", text)
    text = html.unescape(text)
    text = re.sub(r"^[:\s>]+", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    if re.search(r"signed (on|off)|is away|returned at|is idle at", text, re.I):
        return ""
    return text


def parse_conversation(content: str) -> list[tuple[str, str]]:
    entries = []
    matches = list(SPEAKER_RE.finditer(content))
    for i, match in enumerate(matches):
        speaker = norm(match.group(1))
        start = match.end()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(content)
        text = strip_html(content[start:end])
        if text:
            entries.append((speaker, text))
    return entries


def mysql_query(sql: str) -> str:
    return subprocess.check_output(
        ["mysql", "-u", DB["user"], f"-p{DB['password']}", DB["database"], "-N", "-B", "-e", sql],
        text=True,
        errors="replace",
    )


def lookup_account_ids(names: tuple[str, ...]) -> list[int]:
    ids = []
    for name in names:
        q = f"SELECT id FROM accounts WHERE LOWER(REPLACE(screenname,' ',''))=LOWER(REPLACE('{name.replace(chr(39), chr(39)*2)}',' ','')) LIMIT 1"
        out = mysql_query(q).strip()
        if not out:
            raise SystemExit(f"Account not found: {name}")
        ids.append(int(out))
    return ids


def is_good_message(text: str) -> bool:
    if len(text) < 2 or len(text) > 180:
        return False
    if text.startswith("http") or "http://" in text or "www." in text:
        return False
    if re.search(r"[\x00-\x08\x0b\x0c\x0e-\x1f]", text):
        return False
    if re.search(r"\b(TIT2|TYER|WXXX|TENC|LAME)\b", text):
        return False
    if text.count(":") >= 3 and re.search(r"\w+\s*:", text):
        return False
    return True


def voice_stats(msgs: list[str]) -> dict:
    if not msgs:
        return {}
    words = [len(m.split()) for m in msgs]
    slang = Counter()
    for m in msgs:
        for w in re.findall(r"\b[\w']+\b", m.lower()):
            slang[w] += 1
    skip = {
        "the", "a", "an", "to", "and", "or", "i", "you", "it", "is", "in", "of", "that",
        "for", "on", "with", "be", "are", "was", "at", "as", "so", "but", "not", "my", "me",
        "we", "he", "she", "they", "do", "did", "have", "has", "had", "this", "what", "just",
        "like", "get", "got", "all", "if", "can", "would", "will", "your", "u", "ur", "r",
        "0", "d", "http", "com", "www",
    }
    starters = Counter(m.split()[0].lower() for m in msgs if m.split())
    return {
        "message_count": len(msgs),
        "median_words": statistics.median(words),
        "mean_words": round(statistics.mean(words), 2),
        "pct_short_3": round(sum(w <= 3 for w in words) / len(words), 3),
        "pct_short_5": round(sum(w <= 5 for w in words) / len(words), 3),
        "pct_lowercase_start": round(sum(1 for m in msgs if m and m[0].islower()) / len(msgs), 3),
        "pct_no_end_punct": round(
            sum(1 for m in msgs if not re.search(r"[.!?]$", m.rstrip())) / len(msgs), 3
        ),
        "pct_haha": round(sum(1 for m in msgs if re.search(r"haha", m, re.I)) / len(msgs), 3),
        "pct_rofl": round(sum(1 for m in msgs if re.search(r"rofl", m, re.I)) / len(msgs), 3),
        "pct_lol": round(sum(1 for m in msgs if re.search(r"\blol\b", m, re.I)) / len(msgs), 3),
        "pct_ellipsis": round(sum(1 for m in msgs if "..." in m) / len(msgs), 3),
        "pct_caps_word": round(sum(1 for m in msgs if re.search(r"\b[A-Z]{3,}\b", m)) / len(msgs), 3),
        "top_words": [w for w, _ in slang.most_common(60) if w not in skip][:25],
        "starters": [s for s, _ in starters.most_common(20) if s],
    }


def burst_stats(entries: list[tuple[str, str]], persona_key: str) -> dict:
    bursts: list[list[str]] = []
    i = 0
    while i < len(entries):
        if entries[i][0] == persona_key:
            burst = [entries[i][1]]
            j = i + 1
            while j < len(entries) and entries[j][0] == persona_key:
                burst.append(entries[j][1])
                j += 1
            bursts.append(burst)
            i = j
        else:
            i += 1
    lens = [len(b) for b in bursts] or [1]
    return {
        "burst_count": len(bursts),
        "pct_multi": round(sum(l > 1 for l in lens) / len(lens), 3),
        "avg_burst_len": round(statistics.mean(lens), 2),
        "max_burst_len": min(4, max(lens)),
    }


def extract_phrases(msgs: list[str]) -> list[str]:
    phrases = Counter()
    for m in msgs:
        if not is_good_message(m):
            continue
        words = m.split()
        for n in (2, 3):
            for i in range(len(words) - n + 1):
                chunk = " ".join(words[i : i + n]).lower()
                if 4 <= len(chunk) <= 40:
                    phrases[chunk] += 1
    return [p for p, c in phrases.most_common(80) if c >= 3][:40]


def curate_examples(pairs: list[dict], limit: int = 30) -> list[dict]:
    scored = []
    for p in pairs:
        u, a = p["user"], p["assistant"]
        if not is_good_message(u) or not is_good_message(a):
            continue
        if len(a.split()) > 12:
            continue
        score = 0
        if 1 <= len(a.split()) <= 6:
            score += 2
        if "?" in u:
            score += 1
        if re.search(r"\b(sup|hey|haha|wow|doing)\b", u.lower()):
            score += 1
        scored.append((score, p))
    scored.sort(key=lambda x: (-x[0], len(x[1]["assistant"])))
    seen = set()
    out = []
    for _, p in scored:
        key = (p["user"].lower()[:40], p["assistant"].lower()[:40])
        if key in seen:
            continue
        seen.add(key)
        out.append(p)
        if len(out) >= limit:
            break
    return out


def build_pairs(entries: list[tuple[str, str]], persona_key: str, partner: str) -> list[dict]:
    pairs = []
    for i, (speaker, text) in enumerate(entries):
        if speaker == persona_key and i > 0 and entries[i - 1][0] == partner:
            pairs.append({"user": entries[i - 1][1], "assistant": text})
    return pairs


PERSONA_SPEAKER_RE = {
    "magus5311": re.compile(r"(?:Auto response from )?(?:Magus5311|magus5311|MagusisHiding)<!--", re.I),
    "soewaslike": re.compile(r"(?:Auto response from )?(?:soEwaslike|so E was like)<!--", re.I),
}


def extract_appearance(raw_lines: list[str], persona_key: str) -> dict:
    """Pull font/color/bg from raw HTML for this persona's outgoing messages."""
    speaker_re = PERSONA_SPEAKER_RE.get(persona_key)
    if not speaker_re:
        return {
            "font_face": "Trebuchet MS",
            "font_color": "#0000ff",
            "font_size": "1",
            "body_bgcolor": "#000080",
        }

    combos: Counter = Counter()
    backgrounds: Counter = Counter()
    label_colors: Counter = Counter()

    for line in raw_lines:
        if not speaker_re.search(line):
            continue

        bg = re.search(r'BODY BGCOLOR="(#[0-9a-fA-F]+)"', line, re.I)
        if bg:
            backgrounds[bg.group(1).lower()] += 1

        for lm in speaker_re.finditer(line):
            chunk = line[lm.start() : lm.start() + 1200]
            label = re.search(
                r'<FONT COLOR="(#[0-9a-fA-F]+)"[^>]*LANG="0"',
                chunk[:250],
                re.I,
            )
            if label:
                label_colors[label.group(1).lower()] += 1

            rest = chunk[chunk.find("<!--") :]
            rest = rest.split("<!--", 1)[-1]
            rest = rest.split("<!--", 1)[-1] if "<!--" in rest else rest

            for fm in re.finditer(
                r"<FONT\s+([^>]*(?:COLOR=\"#[0-9a-fA-F]+\"[^>]*FACE=\"[^\"]+\"|FACE=\"[^\"]+\"[^>]*COLOR=\"#[0-9a-fA-F]+\"|FACE=\"[^\"]+\")[^>]*)>([^<]{1,200})",
                rest,
                re.I,
            ):
                attrs, text = fm.group(1), fm.group(2).strip()
                if not text or text == ":" or text.startswith("http"):
                    continue
                if re.search(r"signed (on|off)|is away|is idle", text, re.I):
                    continue
                color_m = re.search(r'COLOR="(#[0-9a-fA-F]+)"', attrs, re.I)
                face_m = re.search(r'FACE="([^"]+)"', attrs, re.I)
                size_m = re.search(r"SIZE=(\d+)", attrs, re.I)
                if not face_m:
                    continue
                color = color_m.group(1).lower() if color_m else label_colors.most_common(1)[0][0] if label_colors else "#0000ff"
                combos[(color, face_m.group(1), size_m.group(1) if size_m else "1")] += 1

            for fm in re.finditer(r'<FONT FACE="([^"]+)" SIZE=(\d+)>([^<]{2,200})', rest, re.I):
                face, size, text = fm.group(1), fm.group(2), fm.group(3).strip()
                if text.startswith("http"):
                    continue
                default_color = label_colors.most_common(1)[0][0] if label_colors else "#0000ff"
                combos[(default_color, face, size)] += 1

    defaults = {
        "magus5311": {
            "font_face": "Trebuchet MS",
            "font_color": "#0000ff",
            "font_size": "1",
            "body_bgcolor": "#000040",
            "html_wrap": "full",
        },
        "soewaslike": {
            "font_face": "Lucida Console",
            "font_color": "#208fa4",
            "font_size": "1",
            "body_bgcolor": "#ffffff",
            "html_wrap": "font",
        },
    }
    appearance = defaults.get(persona_key, defaults["magus5311"]).copy()

    if combos:
        (color, face, size), _ = combos.most_common(1)[0]
        appearance["font_color"] = color
        appearance["font_face"] = face
        appearance["font_size"] = size
    if backgrounds:
        appearance["body_bgcolor"] = backgrounds.most_common(1)[0][0]
    if label_colors:
        appearance["label_color"] = label_colors.most_common(1)[0][0]
    appearance.setdefault("html_wrap", appearance.get("html_wrap", "full"))

    return appearance


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("screenname")
    parser.add_argument("--conan-id", type=int, default=10)
    parser.add_argument("--out", default="/opt/aim-bots/personas")
    args = parser.parse_args()

    persona_key = norm(args.screenname)
    alias_names = ACCOUNT_ALIASES.get(persona_key, (persona_key,))
    account_ids = lookup_account_ids(alias_names)
    ids_sql = ",".join(str(i) for i in account_ids)

    raw = mysql_query(
        f"SELECT content FROM conversations WHERE "
        f"(from_account IN ({ids_sql}) AND to_account={args.conan_id}) OR "
        f"(from_account={args.conan_id} AND to_account IN ({ids_sql})) ORDER BY id"
    )
    raw_lines = [line for line in raw.split("\n") if line.strip()]

    raw_all = mysql_query(
        f"SELECT content FROM conversations WHERE "
        f"from_account IN ({ids_sql}) OR to_account IN ({ids_sql}) ORDER BY id"
    )
    raw_all_lines = [line for line in raw_all.split("\n") if line.strip()]

    all_entries: list[tuple[str, str]] = []
    for line in raw_lines:
        all_entries.extend(parse_conversation(line))

    persona_msgs = [text for speaker, text in all_entries if speaker == persona_key]
    pairs = build_pairs(all_entries, persona_key, "conan0fthenight")
    good_samples = list(dict.fromkeys(m for m in persona_msgs if is_good_message(m)))

    appearance = extract_appearance(raw_all_lines, persona_key)

    random.seed(42)
    persona = {
        "screen_name": persona_key,
        "display_name": DISPLAY_NAMES.get(persona_key, args.screenname),
        "merged_accounts": list(alias_names),
        "message_count": len(persona_msgs),
        "pair_count": len(pairs),
        "examples": curate_examples(pairs, 30),
        "sample_messages": random.sample(good_samples, min(60, len(good_samples))),
        "phrases": extract_phrases(persona_msgs),
        **appearance,
    }

    out_dir = Path(args.out) / persona_key
    out_dir.mkdir(parents=True, exist_ok=True)

    voice = VOICE_DEFAULTS.get(persona_key, VOICE_DEFAULTS["magus5311"]).copy()
    voice_path = out_dir / "voice.json"
    if voice_path.exists():
        saved = json.loads(voice_path.read_text())
        stat_keys = {
            "message_count", "median_words", "mean_words", "burst_count", "avg_burst_len",
            "max_burst_len", "top_words", "starters",
        }
        for key, val in saved.items():
            if key not in stat_keys and not key.startswith("pct_"):
                voice[key] = val
    voice.update(voice_stats(persona_msgs))
    voice.update(burst_stats(all_entries, persona_key))

    (out_dir / "persona.json").write_text(json.dumps(persona, indent=2))
    voice_path.write_text(json.dumps(voice, indent=2))
    with gzip.open(out_dir / "pairs.json.gz", "wt", encoding="utf-8") as fh:
        json.dump(pairs, fh)

    print(
        f"Extracted {persona_key} (+{','.join(alias_names)}): "
        f"{len(persona_msgs)} messages, {len(pairs)} pairs, "
        f"median {voice.get('median_words')} words/burst, "
        f"style {appearance.get('font_color')} {appearance.get('font_face')} on {appearance.get('body_bgcolor')}"
    )


if __name__ == "__main__":
    main()