#!/usr/bin/env python3
"""
Sandbox Comparison Runner

Runs all three prompt conditions against the 10 test clusters:
- Condition A: Baseline (current production prompt)
- Condition B: Curiosity Mining (same data, hook-focused prompt)
- Condition C: Unclustered (raw posts, no clustering abstraction)

Saves all results for analysis.
"""

import json
import os
import subprocess
import re
from datetime import datetime

# Import the prompts
from prompt_a_baseline import format_prompt_a
from prompt_b_curiosity_mining import format_prompt_b
from prompt_c_unclustered import format_prompt_c, get_unclustered_posts_sample

SANDBOX_DIR = os.path.dirname(os.path.abspath(__file__))

def call_claude(prompt: str, model: str = "opus", timeout: int = 300) -> dict:
    """Call Claude CLI and return parsed JSON response."""
    try:
        result = subprocess.run(
            ["claude", "--print", "--model", model, "-p", "-"],
            input=prompt,
            capture_output=True,
            text=True,
            timeout=timeout
        )

        if result.returncode != 0:
            return {"error": f"Claude CLI failed: {result.stderr[:200]}"}

        response_text = result.stdout.strip()

        # Parse JSON response
        try:
            # Try to find JSON in the response
            json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
            if json_match:
                return json.loads(json_match.group())
            else:
                return json.loads(response_text)
        except json.JSONDecodeError as e:
            return {
                "error": "JSON parse error",
                "raw_response": response_text[:2000],
                "parse_error": str(e)
            }

    except subprocess.TimeoutExpired:
        return {"error": "Claude CLI timed out"}
    except Exception as e:
        return {"error": str(e)}


def run_condition_a(cluster: dict) -> dict:
    """Run Condition A: Baseline prompt."""
    prompt = format_prompt_a(cluster)
    return call_claude(prompt)


def run_condition_b(cluster: dict) -> dict:
    """Run Condition B: Curiosity mining prompt."""
    prompt = format_prompt_b(cluster)
    return call_claude(prompt)


def run_condition_c(cluster: dict) -> dict:
    """Run Condition C: Unclustered raw posts."""
    # Determine topic area from cluster
    channel = cluster.get('channel', '')
    topic_areas = {
        'brain_is_lying': 'cognitive psychology and behavior',
        'money_traps': 'personal finance and consumer behavior'
    }
    topic_area = topic_areas.get(channel, 'general interest topics')

    # Get raw posts sample
    posts = get_unclustered_posts_sample(cluster)
    prompt = format_prompt_c(posts, topic_area)
    return call_claude(prompt)


def extract_title(result: dict, condition: str) -> str:
    """Extract the title from a result based on condition."""
    if "error" in result:
        return f"[ERROR: {result['error'][:50]}]"

    if condition == "A":
        return result.get("suggested_title", "[No title]")
    elif condition == "B":
        video_angle = result.get("video_angle", {})
        return video_angle.get("title", "[No title]")
    elif condition == "C":
        video_concept = result.get("video_concept", {})
        return video_concept.get("title", "[No title]")
    return "[Unknown]"


def extract_verdict(result: dict, condition: str) -> str:
    """Extract the verdict from a result."""
    if "error" in result:
        return "ERROR"
    return result.get("verdict", "[No verdict]")


def extract_score(result: dict, condition: str) -> float:
    """Extract a numeric score from a result."""
    if "error" in result:
        return 0.0

    if condition == "A":
        return float(result.get("weighted_score", 0))
    elif condition == "B":
        hookability = result.get("hookability_score", {})
        return float(hookability.get("overall", 0))
    elif condition == "C":
        hook_quality = result.get("hook_quality", {})
        return float(hook_quality.get("overall_score", 0))
    return 0.0


def run_full_comparison():
    """Run all conditions on all test clusters."""
    # Load test clusters
    test_path = os.path.join(SANDBOX_DIR, 'test_clusters.json')
    with open(test_path) as f:
        test_data = json.load(f)

    clusters = test_data['clusters']
    print(f"\n{'='*70}")
    print("SANDBOX COMPARISON TEST")
    print(f"{'='*70}")
    print(f"Test clusters: {len(clusters)}")
    print(f"Conditions: A (Baseline), B (Curiosity Mining), C (Unclustered)")
    print(f"Started: {datetime.now().isoformat()}")
    print(f"{'='*70}\n")

    results = []

    for i, cluster in enumerate(clusters):
        test_id = cluster['test_id']
        channel = cluster['channel']
        topic_count = cluster['topic_count']

        print(f"\n[{i+1}/{len(clusters)}] {test_id}")
        print(f"    Topics: {topic_count}, Posts: {len(cluster['raw_posts'])}")
        print(f"    Sample topics: {cluster['sample_topics'][:2]}")

        cluster_result = {
            "test_id": test_id,
            "channel": channel,
            "topic_count": topic_count,
            "sample_topics": cluster['sample_topics'],
            "results": {}
        }

        # Run each condition
        for condition, runner in [("A", run_condition_a), ("B", run_condition_b), ("C", run_condition_c)]:
            print(f"    Running Condition {condition}...", end=" ", flush=True)
            result = runner(cluster)

            title = extract_title(result, condition)
            verdict = extract_verdict(result, condition)
            score = extract_score(result, condition)

            cluster_result["results"][condition] = {
                "title": title,
                "verdict": verdict,
                "score": score,
                "full_response": result
            }

            print(f"→ {verdict} ({score:.1f}) \"{title[:40]}...\"")

        results.append(cluster_result)

        # Save intermediate results after each cluster
        output_path = os.path.join(SANDBOX_DIR, 'comparison_results.json')
        with open(output_path, 'w') as f:
            json.dump({
                "completed": len(results),
                "total": len(clusters),
                "timestamp": datetime.now().isoformat(),
                "results": results
            }, f, indent=2)

    print(f"\n{'='*70}")
    print("COMPARISON COMPLETE")
    print(f"{'='*70}")
    print(f"Results saved to: {output_path}")

    return results


def print_summary(results: list):
    """Print a summary comparison of the results."""
    print(f"\n{'='*70}")
    print("SUMMARY: TITLE COMPARISON")
    print(f"{'='*70}\n")

    for r in results:
        print(f"📊 {r['test_id']}")
        print(f"   Topics: {r['sample_topics'][:2]}")
        print()

        for condition in ["A", "B", "C"]:
            res = r["results"][condition]
            label = {"A": "Baseline", "B": "Curiosity", "C": "Unclustered"}[condition]
            print(f"   [{condition}] {label}: \"{res['title']}\"")
            print(f"       Verdict: {res['verdict']}, Score: {res['score']:.1f}")
        print()
        print("-" * 70)


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "--quick":
        # Quick test: just run on first cluster
        print("Quick test mode: Running on first cluster only")
        with open(os.path.join(SANDBOX_DIR, 'test_clusters.json')) as f:
            test_data = json.load(f)

        cluster = test_data['clusters'][0]
        print(f"\nTest cluster: {cluster['test_id']}")

        print("\nCondition A (Baseline):")
        result_a = run_condition_a(cluster)
        print(json.dumps(result_a, indent=2)[:1000])

        print("\nCondition B (Curiosity Mining):")
        result_b = run_condition_b(cluster)
        print(json.dumps(result_b, indent=2)[:1000])

        print("\nCondition C (Unclustered):")
        result_c = run_condition_c(cluster)
        print(json.dumps(result_c, indent=2)[:1000])
    else:
        # Full run
        results = run_full_comparison()
        print_summary(results)