#!/usr/bin/env python3
"""
Example:

CAI_MODEL="claude-sonnet-4-20250514" CAI_STREAM=True python3 case_study_generator.py --jsonl_file logs/cai_b97af8fc-3d51-45d3-8393-6c3341d33807_20250602_201144_luijait_darwin_24.5.0_81_38_189_27.jsonl --output_php_file alias_web/case_study_test.php

CAI Case Study Generator - Generate PHP case studies from JSONL files.

This script loads context from JSONL files using the same mechanism as CAI's /load command,
runs the UseCase agent with streaming output, and generates PHP case studies.

Usage:
    python case_study_generator.py --jsonl_file logs/session.jsonl --output_php_file output.php
    python case_study_generator.py --jsonl_file logs/last --output_php_file case_studies/latest.php
"""

import os
from dotenv import load_dotenv

# Load .env from current directory only, not from parent directories
dotenv_path = os.path.join(os.getcwd(), '.env')
load_dotenv(dotenv_path=dotenv_path, verbose=False)

# Set default for OPENAI_API_KEY if not already set
if "OPENAI_API_KEY" not in os.environ:
    os.environ["OPENAI_API_KEY"] = ""

import sys
import asyncio
import argparse
from pathlib import Path
import json
import re
from typing import List, Dict, Any, Optional

# Import CAI SDK components
from cai.sdk.agents import Runner
from cai.sdk.agents.models.openai_chatcompletions import message_history, add_to_message_history
from cai.sdk.agents.run_to_jsonl import load_history_from_jsonl
from cai.sdk.agents.stream_events import RunItemStreamEvent
from cai.sdk.agents.items import ToolCallOutputItem

# Import UseCase agent
from src.cai.agents.usecase import use_case_agent

# Rich console for better output
from rich.console import Console
from rich.panel import Panel
from rich.live import Live
from rich.text import Text
from rich.progress import Progress, SpinnerColumn, TextColumn

console = Console()


def extract_php_code(text: str) -> Optional[str]:
    """Extract PHP code from markdown code blocks."""
    if not text:
        return None

    # Try to extract PHP code between ```php and ```
    php_matches = re.findall(r"```php\n(.*?)```", text, re.DOTALL)
    if php_matches:
        return php_matches[0].strip()

    # If no code blocks, check if the entire text looks like PHP
    if text.strip().startswith("<?php") or text.strip().startswith("<!doctype"):
        return text.strip()

    return None


async def generate_case_study(jsonl_file: str, output_php_file: str) -> Optional[str]:
    """
    Generate a PHP case study from a JSONL file using streaming output.

    Args:
        jsonl_file: Path to the JSONL file to load context from
        output_php_file: Path to save the PHP output to

    Returns:
        Path to the saved PHP file or None if failed
    """
    # Clear any existing messages in message_history to start fresh
    message_history.clear()

    # Load context from JSONL file (simulating /load command)
    try:
        console.print(f"[yellow]Loading JSONL file: {jsonl_file}[/yellow]")
        messages = load_history_from_jsonl(jsonl_file)

        if not messages:
            console.print("[red]Error: No messages found in JSONL file[/red]")
            return None

        console.print(f"[green]✓ Loaded {len(messages)} messages from JSONL[/green]")

        # Add messages to message_history (exactly like /load command does)
        for message in messages:
            message_history.append(message)

        # Display loaded context summary
        user_messages = sum(1 for msg in messages if msg.get("role") == "user")
        assistant_messages = sum(1 for msg in messages if msg.get("role") == "assistant")
        tool_messages = sum(1 for msg in messages if msg.get("role") == "tool")

        console.print(
            Panel(
                f"Context loaded:\n"
                f"• User messages: {user_messages}\n"
                f"• Assistant messages: {assistant_messages}\n"
                f"• Tool messages: {tool_messages}",
                title="[bold]JSONL Context Summary[/bold]",
                border_style="blue",
            )
        )

    except Exception as e:
        console.print(f"[red]Error loading JSONL file: {str(e)}[/red]")
        return None

    # Analyze the loaded context to provide better guidance
    context_summary = []
    if messages:
        # Find the main topic/challenge from user messages
        for msg in messages:
            if msg.get("role") == "user" and msg.get("content"):
                content = msg.get("content", "")[:300]  # First 300 chars
                if content and len(content) > 20:  # Skip very short messages
                    context_summary.append(content.strip())
                    if len(context_summary) >= 5:  # Get first few meaningful messages
                        break

    # Generate case study prompt with context
    prompt = "Generate the PHP code for a cybersecurity case study based on the template. "
    prompt += "Analyze the conversation context that has been loaded and create a comprehensive case study. "
    prompt += "Fill in all TEMPLATE-TODO sections with relevant information from the session. "
    prompt += "Explain step by step the problem and the solution in this escenario"
    prompt += "The output should be complete PHP code ready to save to a file."

    # Add a summary of the JSONL conversation to the prompt
    if messages:
        prompt += "\n\n## Conversation Context from JSONL:\n"

        # Get key information from the conversation
        user_msgs = [msg for msg in messages if msg.get("role") == "user"]
        assistant_msgs = [msg for msg in messages if msg.get("role") == "assistant"]
        tool_msgs = [msg for msg in messages if msg.get("role") == "tool"]

        # Add user messages
        if user_msgs:
            prompt += "\n### User Messages:\n"
            for i, msg in enumerate(user_msgs[:5], 1):
                content = msg.get("content", "")[:500]
                if content:
                    prompt += f"{i}. {content}\n"

        # Add key assistant responses
        if assistant_msgs:
            prompt += "\n### Key Assistant Responses:\n"
            for i, msg in enumerate(assistant_msgs[:3], 1):
                content = msg.get("content", "")[:500]
                if content and "I'll help" not in content:  # Skip generic responses
                    prompt += f"{i}. {content}\n"

        # Add tool outputs that might contain important data
        if tool_msgs:
            prompt += "\n### Tool Outputs (key findings):\n"
            important_tools = []
            for msg in tool_msgs:
                content = msg.get("content", "")
                # Look for important patterns in tool output
                if any(
                    keyword in content.lower()
                    for keyword in [
                        "map",
                        "credential",
                        "password",
                        "auth",
                        "endpoint",
                        "192.168",
                        "http",
                    ]
                ):
                    important_tools.append(content[:500])

            for i, content in enumerate(important_tools[:5], 1):
                prompt += f"{i}. {content}\n"

    console.print(f"\n[cyan]Generating case study with UseCase agent...[/cyan]")

    # Configure streaming mode based on environment variable
    stream_mode = os.getenv("CAI_STREAM", "true").lower() != "false"

    try:
        if stream_mode:
            # Streaming mode - similar to CLI implementation
            console.print("[dim]Using streaming mode...[/dim]")

            # Track if we've seen any output
            has_output = False
            accumulated_text = []
            php_code = None

            # Run the streaming process like CLI does
            async def process_streamed_response():
                try:
                    result_stream = Runner.run_streamed(use_case_agent, prompt)

                    with Progress(
                        SpinnerColumn(),
                        TextColumn("[progress.description]{task.description}"),
                        console=console,
                        transient=True,
                    ) as progress:
                        task = progress.add_task(
                            "[cyan]Processing with UseCase agent...", total=None
                        )

                        # Consume events so the async generator is executed
                        async for event in result_stream.stream_events():
                            if isinstance(event, RunItemStreamEvent):
                                # Handle tool outputs
                                if event.name == "tool_output" and isinstance(
                                    event.item, ToolCallOutputItem
                                ):
                                    progress.update(
                                        task,
                                        description=f"[cyan]Tool: {event.item.raw_item.get('name', 'unknown')}...",
                                    )

                                    # Add tool message to history (like CLI does)
                                    tool_msg = {
                                        "role": "tool",
                                        "tool_call_id": event.item.raw_item["call_id"],
                                        "content": event.item.output,
                                    }
                                    add_to_message_history(tool_msg)

                        progress.update(task, description="[green]Finalizing output...")

                    # The result is available after streaming completes
                    # But we need to extract the output from message_history
                    # since streaming doesn't provide direct access to final output

                    # Get the last assistant message from message_history
                    for msg in reversed(message_history):
                        if msg.get("role") == "assistant" and msg.get("content"):
                            return msg.get("content")

                    return None

                except Exception as e:
                    console.print(f"[red]Error in streaming: {str(e)}[/red]")
                    import traceback

                    console.print(f"[red]{traceback.format_exc()}[/red]")
                    return None

            # Run the streaming process
            final_output = await process_streamed_response()

            if final_output:
                php_code = extract_php_code(final_output)
                if not php_code:
                    php_code = final_output

            if php_code:
                console.print(f"[green]✓ Generated {len(php_code)} characters of output[/green]")
            else:
                console.print("[red]Error: No output from UseCase agent[/red]")
                return None

        else:
            # Non-streaming mode (simpler, like in examples)
            console.print("[dim]Using non-streaming mode...[/dim]")

            # Show progress
            with console.status("[bold green]Generating case study...") as status:
                # Instead of passing the conversation history directly,
                # just use the prompt with all the context embedded in it
                # This avoids issues with incomplete tool call/response pairs

                # Run with just the prompt
                result = await Runner.run(use_case_agent, prompt)

            # Extract PHP code from result
            if hasattr(result, "final_output") and result.final_output:
                output_text = result.final_output

                # Process the output to handle tool outputs
                for item in result.new_items:
                    if isinstance(item, ToolCallOutputItem):
                        # Add tool messages to history
                        tool_msg = {
                            "role": "tool",
                            "tool_call_id": item.raw_item["call_id"],
                            "content": item.output,
                        }
                        add_to_message_history(tool_msg)

                php_code = extract_php_code(output_text)
                if not php_code:
                    # If extraction failed, use the raw output
                    php_code = output_text

                console.print(f"[green]✓ Generated {len(php_code)} characters of output[/green]")
            else:
                console.print("[red]Error: No output from UseCase agent[/red]")
                return None

    except Exception as e:
        console.print(f"[red]Error generating case study: {str(e)}[/red]")
        import traceback

        console.print(f"[red]{traceback.format_exc()}[/red]")
        return None

    # Validate PHP code
    if not php_code or len(php_code) < 100:
        console.print("[red]Error: Generated output is too short or invalid[/red]")
        return None

    # Save PHP code to file
    try:
        output_path = Path(output_php_file)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        with open(output_path, "w", encoding="utf-8") as f:
            f.write(php_code)

        console.print(f"\n[green]✓ PHP case study saved to: {output_php_file}[/green]")

        # Display file size and preview
        file_size = output_path.stat().st_size
        console.print(f"[dim]File size: {file_size:,} bytes[/dim]")

        # Show first few lines as preview
        lines = php_code.split("\n")[:15]
        preview = "\n".join(lines)
        if len(php_code.split("\n")) > 15:
            preview += "\n..."

        console.print(Panel(preview, title="[bold]PHP File Preview[/bold]", border_style="blue"))

        return str(output_path)

    except Exception as e:
        console.print(f"[red]Error saving PHP file: {str(e)}[/red]")
        return None


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Generate PHP case studies from JSONL files using CAI UseCase agent.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Generate case study from a specific JSONL file
  python case_study_generator.py --jsonl_file logs/session_20240102_123456.jsonl --output_php_file case_studies/ctf_writeup.php
  
  # Use the last session log (default behavior like /load command)
  python case_study_generator.py --jsonl_file logs/last --output_php_file case_studies/latest.php
  
  # Generate with custom output directory
  python case_study_generator.py --jsonl_file logs/last --output_php_file ~/Documents/case_studies/analysis.php
  
  # Override the model
  python case_study_generator.py --jsonl_file logs/last --output_php_file output.php --model gpt-4o
  
  # Disable streaming
  CAI_STREAM=false python case_study_generator.py --jsonl_file logs/last --output_php_file output.php
        """,
    )
    parser.add_argument(
        "--jsonl_file",
        type=str,
        default="logs/last",
        help="Path to the JSONL file containing conversation context (default: logs/last)",
    )
    parser.add_argument(
        "--output_php_file",
        type=str,
        required=True,
        help="Path where the generated PHP file will be saved",
    )
    parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="Override the model to use (e.g., claude-sonnet-4-20250514, gpt-4o)",
    )
    return parser.parse_args()


async def main():
    """Main entry point for the script."""
    args = parse_args()

    # Display banner
    console.print(
        Panel(
            "[bold cyan]CAI Case Study Generator[/bold cyan]\n"
            "Generate professional cybersecurity case studies from JSONL session logs\n\n"
            "[dim]This tool uses the CAI UseCase agent to analyze session context and generate\n"
            "comprehensive PHP case studies based on the conversation history.[/dim]",
            border_style="cyan",
        )
    )

    # Override model if specified
    if args.model:
        os.environ["CAI_MODEL"] = args.model
        console.print(f"[yellow]Using model override: {args.model}[/yellow]")

    current_model = os.getenv("CAI_MODEL", "alias1")
    console.print(f"[yellow]Model: {current_model}[/yellow]")

    # Check if JSONL file exists
    jsonl_path = Path(args.jsonl_file)
    if not jsonl_path.exists() and args.jsonl_file != "logs/last":
        console.print(f"[red]Error: JSONL file not found: {args.jsonl_file}[/red]")
        return 1

    # Generate the case study
    result = await generate_case_study(args.jsonl_file, args.output_php_file)

    if result:
        console.print("\n[bold green]✨ Case study generation completed successfully![/bold green]")
        console.print(f"[dim]You can now open {result} in your browser or editor[/dim]")
        return 0
    else:
        console.print("\n[bold red]❌ Case study generation failed[/bold red]")
        console.print("[dim]Please check the error messages above and ensure:[/dim]")
        console.print("[dim]1. The JSONL file contains valid session data[/dim]")
        console.print("[dim]2. The UseCase agent has access to the template file[/dim]")
        console.print("[dim]3. Your API keys are properly configured[/dim]")
        return 1


if __name__ == "__main__":
    # Run the async main function
    exit_code = asyncio.run(main())
    sys.exit(exit_code)