"""
import os
from claude_sdk import query, ClaudeCodeOptions, Message, ResultMessage
Here's a simple example of how I might be able to make use of Claude Code in the inner
codegen loop instead of Gemini API directly. I think I'd need to run the output through
another LLM to extract simply the code and nothing else since I can't
"""
import json
from pathlib import Path

import asyncclick as click
from claude_code_sdk import ClaudeCodeOptions, Message, ResultMessage, AssistantMessage, ToolUseBlock, ToolResultBlock, TextBlock, query


def render_todos(todos_input: dict) -> str:
    """Render TodoWrite tool input in a nice format."""
    todos = todos_input.get("todos", [])
    if not todos:
        return "📝 No todos"
    
    result = "📝 Todo List:\n"
    for todo in todos:
        status_emoji = {
            "completed": "✅",
            "in_progress": "🔄", 
            "pending": "⏳"
        }.get(todo["status"], "❓")
        
        priority_emoji = {
            "high": "🔴",
            "medium": "🟡",
            "low": "🟢"
        }.get(todo["priority"], "⚪")
        
        result += f"  {status_emoji} {priority_emoji} {todo['content']}\n"
    
    return result


@click.command()
@click.option(
    '--output_dir', 
    type=click.Path(exists=True, dir_okay=True, file_okay=False, path_type=Path), 
    required=True,
)
@click.option(
    '--desc', 
    required=True,
)
@click.option(
    '--max_turns', 
    type=int, 
    default=10,
)
async def main(output_dir: Path, desc: str, max_turns: int):
    """Generate the high-level scaffolding for a program that's well structured to enable its
    detailed internal implementation can be generated by Jitter.

    Reads the JITTER_DESIGN_AND_USAGE.md file and feeds it to Claude Code along with an explanation
    that it should be generating a program SCAFFOLD that corresponds to whatever the user briefly
    requests.
    """
    JITTER_DESIGN_FILE = "JITTER_DESIGN_AND_USAGE.md"

    try:
        # The design file is expected to be directly next to this file.
        with open(Path(__file__).parent / JITTER_DESIGN_FILE, encoding='utf-8') as f:
            jitter_design_content = f.read()
    except FileNotFoundError:
        print(f"Error: {JITTER_DESIGN_FILE} not found. Please ensure it's in the current directory.")
        return
    except Exception as e:
        print(f"Error reading {JITTER_DESIGN_FILE}: {e}")
        return

    user_program_request = desc 

    if not user_program_request.strip():
        print("Program request cannot be empty. Exiting.")
        return

    system_prompt = (
        "You are an expert program architect specializing in generating high-level code scaffolding. "
        "Your goal is to create a well-structured program outline based on a user's brief request. "
        "The detailed internal implementation of this program will be handled by a tool called Jitter. "
        f"Below is the content of `{JITTER_DESIGN_FILE}`, which describes Jitter's design and usage principles. "
        "Please use this context to ensure the scaffolding you generate is compatible with Jitter's approach, "
        "focusing on clear function signatures, modularity, and placeholders for Jitter-generated components. "
        "Do NOT implement detailed logic; only provide the high-level structure (e.g., type definitions, "
        "function signatures, main execution flow, file structure). "
        "The output should be directly usable Python code with a main entrypoint file in IT'S OWN FILE, do not "
        "put ANY program logic in the file with the entry point. "
        "\n\nIMPORTANT: You MUST actually WRITE the Python files using the Write tool. Do NOT just output "
        "a plan or description - use the Write tool to create the actual .py files in the specified directory. "
        "Create multiple files as needed for proper program structure (main.py, modules, etc.). "
        "Each file should contain working Python code with proper imports and function signatures."
        "\n\n--- JITTER_DESIGN_AND_USAGE.md ---\n"
        f"{jitter_design_content}"
        "\n--- End of JITTER_DESIGN_AND_USAGE.md ---\n"
    )

    user_prompt = f"Generate the high-level scaffolding for the following program: '{user_program_request}'"

    print("\nGenerating scaffolding with Claude Code...")

    messages: list[Message] = []
    turn_count = 0

    try:
        async for message in query(
            prompt=user_prompt,
            options=ClaudeCodeOptions(
                append_system_prompt=system_prompt,
                allowed_tools=[
                    f"Write({output_dir.as_posix()}/*.py)", 
                    f"Edit({output_dir.as_posix()}/*.py)", 
                    f"MultiEdit({output_dir.as_posix()}/*.py)",
                ],
                cwd=output_dir,
                max_turns=max_turns,
                permission_mode='acceptEdits'  # auto-accept file edits
            )
        ):
            messages.append(message)
            if isinstance(message, AssistantMessage):
                turn_count += 1
                print(f"\n--- Turn {turn_count}/{max_turns} ---")
                
                for content_block in message.content:
                    if isinstance(content_block, ToolUseBlock):
                        if content_block.name == "TodoWrite":
                            print(f"\n{render_todos(content_block.input)}")
                        else:
                            print(f"\n🔧 Tool: {content_block.id} - {content_block.name}")
                            print(f"\tInput: {json.dumps(content_block.input, indent=4)}\n\n")
                    elif isinstance(content_block, ToolResultBlock):
                        print(f"\n🔧 Tool Result: {content_block.tool_use_id}")
                        print(f"\tContent: {json.dumps(content_block.content, indent=4)}")
                        print(f"\tError: {content_block.is_error}\n\n.")
                    elif isinstance(content_block, TextBlock):
                        print(content_block.text)
            elif isinstance(message, ResultMessage):
                print("\n" + "*" * 80 + "\nDone!:\n")
                print(message.result)
    except Exception as e:
        print(f"An error occurred during Claude API call: {e}")
        return

if __name__ == "__main__":
    main()
