"""Implementation plan agent for generating technical implementation plans."""

import json
import os
import re
from pathlib import Path
from typing import Optional, List, Tuple

from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt, Confirm

from emdash.agent.toolkit import AgentToolkit
from emdash.agent.runner import SafeJSONEncoder
from emdash.agent.providers import get_provider
from emdash.agent.providers.factory import DEFAULT_MODEL
from emdash.templates import load_template_for_agent
from emdash.agent.spec_schema import Spec


class ImplementationPlanAgent:
    """Agent that generates implementation plans from specifications."""

    MAX_TOOL_RESULT_SIZE = 8000

    def __init__(
        self,
        model: str = DEFAULT_MODEL,
        verbose: bool = True,
        max_iterations: int = 40,
        spec_content: str = "",
        spec_name: str = "",
        project_md_path: Optional[str] = None,
        show_tool_results: bool = False,
    ):
        """Initialize the implementation plan agent.

        Args:
            model: LLM model to use (claude-* for Anthropic, gpt-* for OpenAI)
            verbose: Whether to print progress
            max_iterations: Maximum tool call iterations
            spec_content: The specification content
            spec_name: Name of the specification
            project_md_path: Path to PROJECT.md file
            show_tool_results: Whether to print full tool results (--verbose)
        """
        self.provider = get_provider(model)
        self.toolkit = AgentToolkit(enable_session=True)
        self.model = model
        self.verbose = verbose
        self.show_tool_results = show_tool_results
        self.max_iterations = max_iterations
        self.context_limit = self.provider.get_context_limit()
        self.console = Console()
        self.messages: list[dict] = []
        self.spec_content = spec_content
        self.spec_name = spec_name
        self.project_context = self._load_project_md(project_md_path)

    def _load_project_md(self, path: Optional[str] = None) -> str:
        """Load PROJECT.md if it exists."""
        search_paths = [
            path,
            "PROJECT.md",
            "./PROJECT.md",
            "../PROJECT.md",
        ]

        for p in search_paths:
            if p and os.path.exists(p):
                with open(p, "r") as f:
                    return f.read()
        return ""

    def generate_plan(self) -> str:
        """Generate an implementation plan from the specification.

        Returns:
            The generated implementation plan markdown
        """
        if self.verbose:
            self.console.print(
                Panel(
                    f"[cyan]Generating implementation plan for:[/cyan]\n{self.spec_name}",
                    title="[bold]Implementation Plan Agent[/bold]",
                    border_style="cyan",
                )
            )

        # Build system prompt with project context and spec
        tasks_template = load_template_for_agent("tasks")
        system_content = tasks_template

        # Prepend PROJECT.md if available
        if self.project_context:
            system_content = f"""## PROJECT.md - PROJECT CONTEXT

{self.project_context}

---

{tasks_template}"""
        elif self.verbose:
            self.console.print("[yellow]Warning: No PROJECT.md found.[/yellow]")

        # Append spec
        system_content += f"""

## SPECIFICATION TO IMPLEMENT:

{self.spec_content}
"""

        self.messages = [
            {"role": "system", "content": system_content},
            {"role": "user", "content": """Create an implementation plan for this specification.

MANDATORY SEQUENCE:
1. Call plan_exploration with the feature goal and use_case="implementation"
2. Execute the recommended tools in order - do not skip steps
3. Write a ## Synthesis section summarizing what you learned:
   - Files to modify (paths)
   - Patterns to follow (file:line references)
   - Code owner (from get_file_history commit counts)
   - Blockers/conflicts
4. Generate the implementation plan with single-prompt tasks

Do NOT skip the synthesis step - it ensures you understand the codebase before writing tasks."""},
        ]

        tools = self.toolkit.get_all_schemas()

        # Run exploration and generation loop
        iterations = 0
        while iterations < self.max_iterations:
            iterations += 1

            response = self.provider.chat(self.messages, tools=tools)
            self.messages.append(self.provider.format_assistant_message(response))

            if response.tool_calls:
                # Execute tool calls
                for tool_call in response.tool_calls:
                    result = self._execute_tool_call(tool_call)
                    result_json = json.dumps(result, cls=SafeJSONEncoder)
                    self.messages.append(
                        self.provider.format_tool_result(tool_call.id, result_json)
                    )
            else:
                # Check if response was truncated due to max_tokens
                if response.stop_reason == "max_tokens":
                    if self.verbose:
                        self.console.print("[yellow]Response truncated (max_tokens). Requesting continuation...[/yellow]")
                    self.messages.append({
                        "role": "user",
                        "content": "Your response was cut off. Please continue generating the implementation plan.",
                    })
                    continue

                # Model stopped calling tools - return whatever it generated
                content = response.content or ""
                if content.strip():
                    extracted = self._extract_markdown(content)
                    # Validate it looks like a plan, not tool output
                    if self._is_valid_plan_content(extracted):
                        return extracted
                    elif self.verbose:
                        self.console.print("[yellow]Response doesn't look like a valid plan. Requesting proper format...[/yellow]")

                # Empty or invalid response - ask for the plan
                self.messages.append({
                    "role": "user",
                    "content": "Please generate the implementation plan in markdown format now. Start with '# Implementation Plan:' and include Tasks sections.",
                })

        # Max iterations - generate what we have
        if self.verbose:
            self.console.print("[yellow]Max iterations reached, generating plan...[/yellow]")

        return self._force_plan_generation()

    def _truncate_data(self, data: dict) -> dict:
        """Truncate data to fit within size limits."""
        serialized = json.dumps(data, cls=SafeJSONEncoder)
        if len(serialized) <= self.MAX_TOOL_RESULT_SIZE:
            return data

        truncated = {}
        for key, value in data.items():
            if isinstance(value, list) and len(value) > 10:
                truncated[key] = value[:10]
                truncated[f"{key}_truncated"] = True
                truncated[f"{key}_total"] = len(value)
            elif isinstance(value, dict):
                truncated[key] = self._truncate_data(value)
            else:
                truncated[key] = value

        return truncated

    def _execute_tool_call(self, tool_call) -> dict:
        """Execute a tool call and return the result."""
        name = tool_call.name
        try:
            args = json.loads(tool_call.arguments)
        except json.JSONDecodeError:
            args = {}

        result = self.toolkit.execute(name, **args)

        if self.verbose:
            self._print_tool_call(name, args, result)

        if result.success:
            data = self._truncate_data(result.data)
            return {
                "success": True,
                "data": data,
                "suggestions": result.suggestions,
            }
        else:
            return {
                "success": False,
                "error": result.error,
                "suggestions": result.suggestions,
            }

    def _print_tool_call(self, name: str, args: dict, result):
        """Print concise tool call info."""
        status = "[green]✓[/green]" if result.success else "[red]✗[/red]"

        args_str = ""
        if args:
            key_args = []
            for k, v in list(args.items())[:2]:
                if isinstance(v, str) and len(v) > 30:
                    v = v[:30] + "..."
                key_args.append(f"{k}={v}")
            args_str = f" ({', '.join(key_args)})"

        result_str = ""
        if result.success and result.data:
            if "results" in result.data:
                result_str = f" → {len(result.data['results'])} results"
            elif "commits" in result.data:
                result_str = f" → {len(result.data['commits'])} commits"
            elif "files" in result.data:
                result_str = f" → {len(result.data['files'])} files"

        self.console.print(f"  {status} [cyan]{name}[/cyan]{args_str}{result_str}")

        # Print full results if --verbose flag is set
        if self.show_tool_results and result.success and result.data:
            self.console.print()
            self.console.print(f"  [dim]─── {name} result ───[/dim]")
            result_json = json.dumps(result.data, indent=2, default=str)
            # Truncate very long results
            if len(result_json) > 3000:
                result_json = result_json[:3000] + "\n... (truncated)"
            self.console.print(f"  [dim]{result_json}[/dim]")
            self.console.print()

    def _force_plan_generation(self) -> str:
        """Request the model to generate a plan."""
        if self.verbose:
            self.console.print("[dim]Requesting implementation plan...[/dim]")

        self.messages.append({
            "role": "user",
            "content": "Generate the implementation plan now based on what you've learned.",
        })

        response = self.provider.chat(self.messages)
        content = response.content or ""

        return self._extract_markdown(content)

    def _extract_markdown(self, content: str) -> str:
        """Extract markdown content from response."""
        if "```markdown" in content:
            start = content.find("```markdown") + len("```markdown")
            end = content.find("```", start)
            if end > start:
                return content[start:end].strip()

        if "```" in content and "# Implementation Plan:" in content:
            start = content.find("```") + 3
            end = content.find("```", start)
            if end > start:
                return content[start:end].strip()

        if "# Implementation Plan:" in content:
            return content.strip()

        return content

    def _is_valid_plan_content(self, content: str) -> bool:
        """Check if content looks like a valid implementation plan."""
        if not content or len(content) < 100:
            return False

        # Must have markdown headers
        has_headers = "#" in content

        # Should have plan-like keywords
        plan_keywords = ["task", "step", "file", "implement", "plan"]
        has_plan_keywords = any(kw in content.lower() for kw in plan_keywords)

        # Should NOT look like raw JSON tool output
        looks_like_json = content.strip().startswith("{") or content.strip().startswith("[")
        has_tool_fields = '"caller"' in content or '"callee"' in content or '"qualified_name"' in content

        if looks_like_json or has_tool_fields:
            return False

        return has_headers and has_plan_keywords


def list_specs(specs_dir: str = "specs") -> List[Tuple[str, str]]:
    """List all available specifications.

    Args:
        specs_dir: Directory containing specs

    Returns:
        List of (name, path) tuples
    """
    specs = []
    specs_path = Path(specs_dir)

    if not specs_path.exists():
        return specs

    for spec_dir in specs_path.iterdir():
        if spec_dir.is_dir():
            spec_file = spec_dir / "spec.json"
            if spec_file.exists():
                specs.append((spec_dir.name, str(spec_file)))

    return sorted(specs)


def find_spec(name: str, specs_dir: str = "specs") -> Optional[Tuple[str, str]]:
    """Find a spec by name (partial match).

    Args:
        name: Full or partial spec name
        specs_dir: Directory containing specs

    Returns:
        (name, path) tuple if found, None otherwise
    """
    specs = list_specs(specs_dir)
    name_lower = name.lower()

    # Exact match first
    for spec_name, spec_path in specs:
        if spec_name.lower() == name_lower:
            return (spec_name, spec_path)

    # Partial match
    matches = []
    for spec_name, spec_path in specs:
        if name_lower in spec_name.lower():
            matches.append((spec_name, spec_path))

    if len(matches) == 1:
        return matches[0]
    elif len(matches) > 1:
        return None  # Ambiguous, caller should show options

    return None


def load_spec(spec_path: str) -> str:
    """Load specification content from file.

    Args:
        spec_path: Path to spec file (markdown or JSON)

    Returns:
        Specification content as markdown
    """
    with open(spec_path, "r") as f:
        content = f.read()

    # If it's a markdown file, return as-is
    if spec_path.endswith(".md"):
        return content

    # Try to parse as JSON (legacy format)
    try:
        data = json.loads(content)
        # Convert legacy JSON to markdown
        title = data.get("title", "Untitled Spec")
        spec = Spec(title=title, content=json.dumps(data, indent=2))
        return spec.to_markdown()
    except json.JSONDecodeError:
        # Assume it's markdown
        return content
