# %% [markdown]
# # Notebook 7: Prompt Complexity & Specificity Scoring
# Score each user message, correlate with session length, tool count, and success.

# %%
import sys, os, re
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import pandas as pd
import numpy as np
from utils.connection import init, query_df
from utils import sessions
from IPython.display import display, Markdown

pd.set_option("display.max_columns", 30)
pd.set_option("display.max_rows", 100)
pd.set_option("display.max_colwidth", 120)

conn = init()
display(Markdown("**Connected.**"))

# %% [markdown]
# ## 1 — Load user messages

# %%
user_msgs = query_df(conn, """
    SELECT m.id, m.session_id, m.content, m.timestamp, m.source,
           s.user_email, s.org, s.repo_name
    FROM   messages m
    JOIN   sessions s ON s.id = m.session_id
    WHERE  m.msg_type = 'user'
    ORDER  BY m.timestamp ASC
""")

# Strip instruction tags — these are Codex system prompts, not user content
def strip_tags(text):
    if not text:
        return ""
    text = re.sub(r"<INSTRUCTIONS?>.*?</INSTRUCTIONS?>", "", text, flags=re.DOTALL | re.IGNORECASE)
    return text.strip()

user_msgs["clean"] = user_msgs["content"].apply(strip_tags)
user_msgs["word_count"] = user_msgs["clean"].str.split().str.len().fillna(0).astype(int)

display(Markdown(f"**User messages:** {len(user_msgs)}"))

# %% [markdown]
# ## 2 — Complexity score (0-6)

# %%
def complexity_score(text):
    """Higher = more complex prompt."""
    if not text:
        return 0
    score = 0
    wc = len(text.split())
    # Length
    if wc > 50:
        score += 1
    if wc > 150:
        score += 1
    # Bullet points / numbered lists
    if re.search(r"^\s*[-*]\s|\d+\.\s", text, re.MULTILINE):
        score += 1
    # Multiple file references
    file_refs = re.findall(r"[\w/.-]+\.\w{1,6}", text)
    if len(file_refs) >= 2:
        score += 1
    # Code blocks
    if "```" in text:
        score += 1
    # Explicit constraints
    if re.search(r"\b(must|should not|without breaking|don't|do not|ensure|make sure)\b", text, re.IGNORECASE):
        score += 1
    return score

user_msgs["complexity"] = user_msgs["clean"].apply(complexity_score)

display(Markdown("### Complexity score distribution"))
display(user_msgs["complexity"].value_counts().sort_index().to_frame("count"))

# %% [markdown]
# ## 3 — Specificity score (0-5)

# %%
def specificity_score(text):
    """Higher = more specific prompt."""
    if not text:
        return 0
    score = 0
    if re.search(r"[\w/]+\.\w{1,5}\b", text):
        score += 1
    if re.search(r"\b[a-z]+[A-Z]\w+\b|\b[a-z]+_[a-z]+\w*\b|\b[A-Z][a-z]+[A-Z]\w+\b", text):
        score += 1
    if re.search(r"(Error|Exception|Traceback|TypeError|ValueError|stack trace)", text, re.IGNORECASE):
        score += 1
    if re.search(r"\bline\s+\d+\b|\bL\d+\b", text, re.IGNORECASE):
        score += 1
    if "```" in text or re.search(r"^\s{4,}\S", text, re.MULTILINE):
        score += 1
    return score

user_msgs["specificity"] = user_msgs["clean"].apply(specificity_score)

display(Markdown("### Specificity score distribution"))
display(user_msgs["specificity"].value_counts().sort_index().to_frame("count"))

# %% [markdown]
# ## 4 — Per-user complexity vs specificity

# %%
display(Markdown("### Avg scores per user"))
display(user_msgs.groupby("user_email")[["complexity", "specificity", "word_count"]].mean().round(2))

# %% [markdown]
# ## 5 — 2D scatter data: complexity vs specificity

# %%
# Aggregate per session (first message or avg of all)
first_msgs = user_msgs.sort_values("timestamp").groupby("session_id").first().reset_index()

display(Markdown("### First-message complexity vs specificity (per session)"))
display(first_msgs.groupby("user_email")[["complexity", "specificity"]].describe().round(2))

# Cross-tab
display(Markdown("### Complexity × Specificity cross-tab (first messages)"))
display(pd.crosstab(first_msgs["complexity"], first_msgs["specificity"], margins=True))

# %% [markdown]
# ## 6 — Correlate with session metrics

# %%
# Load session metrics
sess_metrics = query_df(conn, """
    SELECT m.session_id,
           COUNT(*) AS total_msgs,
           COUNT(*) FILTER (WHERE m.msg_type = 'tool_call') AS tool_calls,
           COUNT(*) FILTER (WHERE m.msg_type = 'user') AS user_msgs
    FROM   messages m
    GROUP  BY m.session_id
""")

all_sess = sessions.list_all(conn, limit=500)
all_sess["duration_min"] = (
    (all_sess["last_updated"] - all_sess["first_seen"]).dt.total_seconds() / 60
)

first_msgs = first_msgs.merge(sess_metrics, on="session_id", how="left")
first_msgs = first_msgs.merge(
    all_sess[["id", "duration_min"]].rename(columns={"id": "session_id"}),
    on="session_id", how="left"
)

display(Markdown("### Correlation: first-message complexity vs session metrics"))
corr_cols = ["complexity", "specificity", "word_count", "total_msgs", "tool_calls", "duration_min"]
available_cols = [c for c in corr_cols if c in first_msgs.columns]
display(first_msgs[available_cols].corr().round(3))

# %% [markdown]
# ## 7 — Key question: do specific prompts lead to shorter sessions?

# %%
display(Markdown("### Session duration by first-message specificity"))
display(first_msgs.groupby("specificity")[["duration_min", "tool_calls", "total_msgs"]].mean().round(1))

display(Markdown("### Session duration by first-message complexity"))
display(first_msgs.groupby("complexity")[["duration_min", "tool_calls", "total_msgs"]].mean().round(1))

# %% [markdown]
# ## 8 — Vague prompts → more exploration?

# %%
first_msgs["is_vague"] = (first_msgs["word_count"] < 10) & (first_msgs["specificity"] == 0)

display(Markdown("### Vague first-message vs non-vague: session comparison"))
display(first_msgs.groupby("is_vague")[["duration_min", "tool_calls", "total_msgs", "user_msgs"]].mean().round(1))

# %% [markdown]
# ## 9 — Binned analysis: low/medium/high specificity

# %%
first_msgs["spec_bucket"] = pd.cut(first_msgs["specificity"], bins=[-1, 0, 1, 5], labels=["none", "low", "high"])

display(Markdown("### Session metrics by specificity bucket"))
display(first_msgs.groupby("spec_bucket")[["duration_min", "tool_calls", "total_msgs"]].agg(["mean", "median", "count"]).round(1))

# %% [markdown]
# ---
# *End of Notebook 7.*
