prompt_parser.py

"""
TimeLapseForge — Prompt Parser Module
Parses GPT JSON output and extracts panel prompts, style guides, and assembly instructions.
Also provides a quick-generate template for text-only input.
"""

import json
import re
from typing import List, Dict, Any, Optional, Tuple


class PromptParser:
    """Parses and validates timelapse JSON from GPT output."""

    @staticmethod
    def clean_json_string(raw: str) -> str:
        """Remove markdown code block wrappers and clean the string."""
        raw = raw.strip()
        # Remove ```json ... ``` or ``` ... ```
        pattern = r"^```(?:json)?\s*\n?(.*?)\n?\s*```$"
        match = re.match(pattern, raw, re.DOTALL)
        if match:
            raw = match.group(1).strip()
        return raw

    @staticmethod
    def find_json_in_text(text: str) -> str:
        """Find the first valid JSON object in a block of text."""
        # Try to find JSON between curly braces
        depth = 0
        start = None
        for i, char in enumerate(text):
            if char == '{':
                if depth == 0:
                    start = i
                depth += 1
            elif char == '}':
                depth -= 1
                if depth == 0 and start is not None:
                    return text[start:i + 1]
        return text

    def parse(self, raw_input: str) -> Dict[str, Any]:
        """
        Parse raw JSON string from GPT output.
        Handles markdown wrappers, nested JSON, and common formatting issues.
        """
        cleaned = self.clean_json_string(raw_input)
        
        try:
            data = json.loads(cleaned)
            return {"success": True, "data": data, "error": None}
        except json.JSONDecodeError:
            pass

        # Try finding JSON within text
        try:
            json_str = self.find_json_in_text(cleaned)
            data = json.loads(json_str)
            return {"success": True, "data": data, "error": None}
        except (json.JSONDecodeError, Exception) as e:
            return {"success": False, "data": None, "error": str(e)}

    def extract_prompts(self, data: Dict[str, Any]) -> List[Dict[str, str]]:
        """Extract all panel prompts from parsed JSON data."""
        panels = data.get("panels", [])
        prompts = []

        for panel in panels:
            panel_id = panel.get("panel_id", len(prompts) + 1)
            
            # Handle different JSON structures
            img_prompt = panel.get("image_prompt", {})
            if isinstance(img_prompt, dict):
                main_prompt = img_prompt.get("main_prompt", "")
                negative = img_prompt.get("negative_prompt", "")
                style = img_prompt.get("style_suffix", "")
            elif isinstance(img_prompt, str):
                main_prompt = img_prompt
                negative = ""
                style = ""
            else:
                # Fallback: look for prompt directly in panel
                main_prompt = panel.get("prompt", panel.get("description", ""))
                negative = panel.get("negative_prompt", "")
                style = ""

            # Also check for video_prompt
            vid_prompt = panel.get("video_prompt", {})
            motion_desc = ""
            if isinstance(vid_prompt, dict):
                motion_desc = vid_prompt.get("motion_description", "")

            prompts.append({
                "panel_id": panel_id,
                "main_prompt": main_prompt,
                "negative_prompt": negative,
                "style_suffix": style,
                "motion_description": motion_desc,
                "panel_title": panel.get("panel_title", f"Panel {panel_id}"),
                "timestamp_label": panel.get("timestamp_label", ""),
                "phase": panel.get("phase_name", panel.get("phase", "")),
            })

        return prompts

    def extract_style_guide(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Extract style guide settings."""
        return data.get("style_guide", {})

    def extract_assembly_guide(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Extract video assembly guide."""
        return data.get("video_assembly_guide", {})

    def extract_metadata(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Extract project metadata."""
        return data.get("project_metadata", {})

    def get_summary(self, data: Dict[str, Any]) -> str:
        """Get a human-readable summary of the parsed JSON."""
        meta = self.extract_metadata(data)
        prompts = self.extract_prompts(data)
        phases = data.get("phases", data.get("restoration_phases", []))

        summary = f"📋 **Project:** {meta.get('project_title', 'Untitled')}\n"
        summary += f"🎬 **Type:** {meta.get('project_type', 'Unknown')}\n"
        summary += f"🖼️ **Total Panels:** {len(prompts)}\n"
        summary += f"📊 **Phases:** {len(phases)}\n"
        summary += f"⏱️ **Timespan:** {meta.get('estimated_real_world_timespan', 'Unknown')}\n\n"

        summary += "**Panels Preview:**\n"
        for p in prompts[:5]:
            summary += f"- Panel {p['panel_id']}: {p['main_prompt'][:80]}...\n"
        if len(prompts) > 5:
            summary += f"- ... and {len(prompts) - 5} more panels\n"

        return summary


class QuickGenerator:
    """Generate basic panel JSON from a simple text description."""

    RESTORATION_PHASES = [
        ("Initial State", 0.0, 0.05, "completely damaged, worn, broken, dirty, neglected"),
        ("Assessment & Cleanup", 0.05, 0.15, "being assessed, initial cleaning, removing loose debris"),
        ("Disassembly & Stripping", 0.15, 0.25, "parts being removed, old paint stripped, exposing bare material"),
        ("Repair & Structural Work", 0.25, 0.45, "active repair, welding, filling dents, replacing broken parts"),
        ("Sanding & Preparation", 0.45, 0.55, "surface sanding, smoothing, preparing for primer"),
        ("Priming", 0.55, 0.65, "primer coat applied, even grey surface, smooth"),
        ("Painting", 0.65, 0.78, "paint being applied, color emerging, wet paint sheen"),
        ("Clear Coat & Finishing", 0.78, 0.88, "clear coat applied, glossy finish, protective layer"),
        ("Reassembly & Detailing", 0.88, 0.95, "parts reinstalled, chrome polished, details perfected"),
        ("Final Reveal", 0.95, 1.0, "fully restored, pristine, gleaming, showroom quality, dramatic lighting"),
    ]

    CREATION_PHASES = [
        ("Empty Start", 0.0, 0.05, "empty space, blank canvas, raw materials, nothing built yet"),
        ("Foundation & Planning", 0.05, 0.15, "ground preparation, foundation laid, basic framework started"),
        ("Core Structure", 0.15, 0.35, "main structure being built, walls rising, skeleton forming"),
        ("Enclosure & Walls", 0.35, 0.50, "walls completed, roof structure, enclosed space taking shape"),
        ("Systems & Infrastructure", 0.50, 0.62, "wiring, plumbing, mechanical systems being installed"),
        ("Interior Rough Work", 0.62, 0.73, "insulation, drywall, rough interior taking shape"),
        ("Finishing & Details", 0.73, 0.85, "paint, trim, fixtures, flooring being installed"),
        ("Furnishing & Styling", 0.85, 0.95, "furniture placed, decorations added, styled and arranged"),
        ("Completed Reveal", 0.95, 1.0, "fully completed, beautiful, perfect lighting, hero shot"),
    ]

    def generate(
        self,
        description: str,
        num_panels: int = 20,
        mode: str = "restoration",
        style: str = "photorealistic, 4K, cinematic lighting, consistent camera angle, shot on Sony A7IV"
    ) -> Dict[str, Any]:
        """Generate a complete JSON structure from a text description."""

        phases = self.RESTORATION_PHASES if mode == "restoration" else self.CREATION_PHASES
        panels = []

        for i in range(num_panels):
            progress = i / max(num_panels - 1, 1)

            # Find current phase
            current_phase = phases[-1]
            for phase in phases:
                if phase[1] <= progress <= phase[2]:
                    current_phase = phase
                    break

            phase_name, _, _, phase_desc = current_phase
            progress_pct = progress * 100

            main_prompt = (
                f"{description}, {phase_desc}, "
                f"restoration progress {progress_pct:.0f}% complete, "
                f"detailed textures visible, consistent environment, "
                f"same camera angle throughout, {style}"
            )

            negative = (
                "blurry, low quality, cartoon, anime, drawing, sketch, text, watermark, "
                "inconsistent angle, different background, teleportation, sudden changes, "
                "extra limbs, deformed, ugly, duplicate"
            )

            panels.append({
                "panel_id": i + 1,
                "phase_name": phase_name,
                "panel_title": f"{phase_name} — {progress_pct:.0f}%",
                "timestamp_label": f"Step {i + 1}/{num_panels}",
                "image_prompt": {
                    "main_prompt": main_prompt,
                    "negative_prompt": negative,
                    "style_suffix": style
                },
                "video_prompt": {
                    "motion_description": f"Subtle ambient motion, {phase_desc}",
                    "camera_motion": "static, locked tripod"
                }
            })

        return {
            "project_metadata": {
                "project_title": f"Timelapse: {description[:50]}",
                "project_type": mode.upper(),
                "total_panels": num_panels,
                "user_original_command": description
            },
            "style_guide": {
                "art_style": "photorealistic",
                "camera": {"type": "fixed-tripod", "focal_length": "35mm"},
                "lighting": {"primary_source": "natural daylight"}
            },
            "panels": panels,
            "video_assembly_guide": {
                "recommended_fps": 24,
                "frame_hold_duration": 2,
                "total_estimated_duration": f"{num_panels * 2} seconds"
            }
        }