#!/usr/bin/env python3
"""
Validate microdrama script for logical consistency using Claude.

Uses LLM to semantically detect:
1. Object timeline - objects can't appear before their introduction scene
2. Character descriptions - shouldn't contain scene-specific objects
3. Frame descriptions - objects mentioned should be in scene's objects list

Usage:
    python validate_script.py /path/to/task_dir
    python validate_script.py /path/to/script.json
"""

import json
import os
import sys
from pathlib import Path
from typing import Dict, List, Tuple

import requests
from dotenv import load_dotenv

# Load .env from parent directory
env_path = Path(__file__).parent.parent.parent / ".env"
load_dotenv(env_path)

LAOZHANG_KEY = os.getenv("LAOZHANG_KEY")
CLAUDE_MODEL = "claude-sonnet-4-20250514"


def call_claude(prompt: str, max_tokens: int = 2000) -> str:
    """Call Claude via LaoZhang API."""
    response = requests.post(
        "https://api.laozhang.ai/v1/messages",
        headers={
            "x-api-key": LAOZHANG_KEY,
            "Content-Type": "application/json",
            "anthropic-version": "2023-06-01"
        },
        json={
            "model": CLAUDE_MODEL,
            "max_tokens": max_tokens,
            "messages": [{"role": "user", "content": prompt}]
        },
        timeout=60
    )

    if response.status_code != 200:
        print(f"Claude API error: {response.status_code}")
        print(response.text[:500])
        return ""

    result = response.json()
    return result.get("content", [{}])[0].get("text", "")


class ScriptValidator:
    def __init__(self, script: dict):
        self.script = script
        self.errors: List[str] = []
        self.warnings: List[str] = []

        # Extract definitions
        self.characters = script.get("characters", {})
        self.objects = script.get("objects", {})
        self.locations = script.get("locations", {})
        self.storyteller = script.get("storyteller", {})
        self.scenes = script.get("scenes", [])

    def _get_object_first_appearance(self) -> Dict[str, int]:
        """Determine when each object first appears in scenes."""
        first_appearance = {}
        for scene in self.scenes:
            scene_id = scene.get("scene_id", 0)
            scene_objects = scene.get("objects", [])
            for obj_id in scene_objects:
                if obj_id not in first_appearance:
                    first_appearance[obj_id] = scene_id
        return first_appearance

    def validate_with_claude(self):
        """Use Claude to semantically validate object references in scenes."""
        if not LAOZHANG_KEY:
            print("Warning: LAOZHANG_KEY not set, skipping Claude validation")
            return

        first_appearance = self._get_object_first_appearance()

        # Build object descriptions for Claude
        objects_desc = "\n".join([
            f"- {obj_id}: {obj_data.get('description', '')}"
            for obj_id, obj_data in self.objects.items()
        ])

        # Build scenes summary
        scenes_summary = []
        for scene in self.scenes:
            scene_id = scene.get("scene_id", "?")
            scene_objects = scene.get("objects", [])
            frames = scene.get("frames", {})

            scenes_summary.append(f"""
Scene {scene_id}:
  Declared objects: {scene_objects if scene_objects else "none"}
  First frame: {frames.get('first', '')}
  Last frame: {frames.get('last', '')}""")

        # Build character descriptions
        all_chars = dict(self.characters)
        if self.storyteller:
            all_chars["storyteller"] = self.storyteller

        chars_desc = "\n".join([
            f"- {char_id}: {char_data.get('description', '')}"
            for char_id, char_data in all_chars.items()
        ])

        prompt = f"""You are a script validator for a video generation pipeline.

OBJECTS DEFINED IN THIS SCRIPT:
{objects_desc}

OBJECT TIMELINE (when each object first appears):
{json.dumps(first_appearance, indent=2)}

CHARACTER DESCRIPTIONS:
{chars_desc}

SCENES:
{''.join(scenes_summary)}

YOUR TASK:
Analyze each scene's frame descriptions and check for logical consistency:

1. OBJECT REFERENCES IN FRAMES: If a frame description mentions or implies an object (the book, money, banknote, etc.), check:
   - Is that object in the scene's "Declared objects" list?
   - Has the object been introduced yet (check object timeline)?

   Note: Generic furniture like "bookshelf" is NOT the same as "the book object".
   Only flag when the SPECIFIC object from the objects list is clearly referenced.

2. CHARACTER DESCRIPTIONS: Check if any character description contains references to
   scene-specific objects that should only appear in certain scenes.

Output your findings as JSON:
{{
  "errors": [
    {{"scene_id": N, "frame": "first|last", "issue": "description of the problem"}},
    ...
  ],
  "warnings": [
    {{"type": "warning_type", "details": "description"}},
    ...
  ],
  "summary": "Brief validation summary"
}}

Be strict but not overly pedantic. Generic references to furniture or environment are OK.
Focus on the SPECIFIC tracked objects that appear in the story.
"""

        print("Calling Claude for semantic validation...")
        response = call_claude(prompt)

        if not response:
            print("Warning: No response from Claude")
            return

        # Parse JSON response
        try:
            # Extract JSON from response (might have markdown)
            json_start = response.find("{")
            json_end = response.rfind("}") + 1
            if json_start >= 0 and json_end > json_start:
                result = json.loads(response[json_start:json_end])

                for error in result.get("errors", []):
                    scene_id = error.get("scene_id", "?")
                    frame = error.get("frame", "")
                    issue = error.get("issue", "")
                    self.errors.append(f"Scene {scene_id} {frame} frame: {issue}")

                for warning in result.get("warnings", []):
                    self.warnings.append(f"{warning.get('type', '')}: {warning.get('details', '')}")

                if result.get("summary"):
                    print(f"\n📋 Claude's summary: {result['summary']}\n")
        except json.JSONDecodeError as e:
            print(f"Warning: Could not parse Claude response as JSON: {e}")
            print(f"Response: {response[:500]}...")

    def validate_scene_completeness(self):
        """Check that scenes have required fields."""
        for scene in self.scenes:
            scene_id = scene.get("scene_id", "?")

            # Check for location (only warn, flashbacks might have minimal setup)
            if not scene.get("locations"):
                self.warnings.append(
                    f"Scene {scene_id} has no locations specified"
                )

            # Check frames exist
            frames = scene.get("frames", {})
            if not frames.get("first"):
                self.errors.append(f"Scene {scene_id} missing 'frames.first' description")
            if not frames.get("last"):
                self.errors.append(f"Scene {scene_id} missing 'frames.last' description")

            # Check captions
            if not scene.get("caption_first"):
                self.warnings.append(f"Scene {scene_id} missing caption_first")
            if not scene.get("caption_last"):
                self.warnings.append(f"Scene {scene_id} missing caption_last")

    def validate_references_exist(self):
        """Check that all referenced characters/objects/locations are defined."""
        for scene in self.scenes:
            scene_id = scene.get("scene_id", "?")

            # Check characters
            for char_id in scene.get("characters", []):
                if char_id not in self.characters and char_id != "storyteller":
                    self.errors.append(
                        f"Scene {scene_id} references undefined character '{char_id}'"
                    )

            # Check objects
            for obj_id in scene.get("objects", []):
                if obj_id not in self.objects:
                    self.errors.append(
                        f"Scene {scene_id} references undefined object '{obj_id}'"
                    )

            # Check locations
            for loc_id in scene.get("locations", []):
                if loc_id not in self.locations:
                    self.errors.append(
                        f"Scene {scene_id} references undefined location '{loc_id}'"
                    )

    def validate_sheets_specified(self):
        """Check that all characters/objects/locations have sheet files specified."""
        all_items = [
            ("character", self.characters),
            ("object", self.objects),
            ("location", self.locations),
        ]

        if self.storyteller:
            all_items.append(("storyteller", {"storyteller": self.storyteller}))

        for item_type, items in all_items:
            for item_id, item_data in items.items():
                sheet = item_data.get("sheet") or item_data.get("reference")
                if not sheet:
                    self.warnings.append(
                        f"{item_type.capitalize()} '{item_id}' has no sheet/reference file specified"
                    )

    def validate(self) -> Tuple[List[str], List[str]]:
        """Run all validations and return (errors, warnings)."""
        self.errors = []
        self.warnings = []

        # Basic structural validation
        self.validate_scene_completeness()
        self.validate_references_exist()
        self.validate_sheets_specified()

        # Semantic validation with Claude
        self.validate_with_claude()

        return self.errors, self.warnings

    def print_object_timeline(self):
        """Print object timeline for debugging."""
        first_appearance = self._get_object_first_appearance()

        print("\n📦 OBJECT TIMELINE:")
        print("-" * 40)
        for obj_id in self.objects:
            first_scene = first_appearance.get(obj_id, "never introduced in scenes")
            desc = self.objects[obj_id].get("description", "")[:50]
            print(f"  {obj_id}:")
            print(f"    First appears: scene {first_scene}")
            print(f"    Description: {desc}...")
        print()


def main():
    if len(sys.argv) < 2:
        print("Usage: python validate_script.py <task_dir_or_script.json>")
        sys.exit(1)

    path = Path(sys.argv[1])

    # Find script file
    if path.is_dir():
        # Look for script in task directory
        for version in ["script.json", "script_v6.json", "script_v5.json", "script_v4.json", "script_v3.json"]:
            script_path = path / "script" / version
            if script_path.exists():
                break
        else:
            print(f"Error: No script found in {path / 'script'}")
            sys.exit(1)
    else:
        script_path = path

    print(f"Validating: {script_path}")
    print("=" * 60)

    with open(script_path) as f:
        script = json.load(f)

    validator = ScriptValidator(script)
    errors, warnings = validator.validate()

    # Print object timeline
    validator.print_object_timeline()

    # Print results
    if errors:
        print("❌ ERRORS:")
        print("-" * 40)
        for i, error in enumerate(errors, 1):
            print(f"{i}. {error}")
        print()

    if warnings:
        print("⚠️  WARNINGS:")
        print("-" * 40)
        for i, warning in enumerate(warnings, 1):
            print(f"{i}. {warning}")
        print()

    if not errors and not warnings:
        print("✅ Script validation passed!")
    elif not errors:
        print(f"✅ No errors, {len(warnings)} warnings")
    else:
        print(f"❌ {len(errors)} errors, {len(warnings)} warnings")
        sys.exit(1)


if __name__ == "__main__":
    main()
