"""Code generator using pluggable LLM backends with structured JSON output."""

from __future__ import annotations

import ast
import os
import sys
import time
import random
import json
from dataclasses import dataclass
from typing import Optional

from scry_run.console import warning


@dataclass
class GeneratedCode:
    """Result of code generation."""

    code: str
    code_type: str  # method, property, classmethod, staticmethod
    docstring: str
    dependencies: list[str]
    packages: list[str] = None  # PyPI packages to install

    def __post_init__(self):
        if self.packages is None:
            self.packages = []


# JSON schema for structured output
CODE_GENERATION_SCHEMA = {
    "type": "object",
    "properties": {
        "code": {
            "type": "string",
            "description": "The Python code for the requested method/property. Must be valid Python that can be exec'd."
        },
        "code_type": {
            "type": "string",
            "enum": ["method", "property", "classmethod", "staticmethod"],
            "description": "The type of code being generated"
        },
        "docstring": {
            "type": "string",
            "description": "A brief description of what this code does"
        },
        "dependencies": {
            "type": "array",
            "items": {"type": "string"},
            "description": "List of import statements needed (e.g., 'import json', 'from pathlib import Path')"
        },
        "packages": {
            "type": "array",
            "items": {"type": "string"},
            "description": "List of PyPI packages to install (e.g., 'pygame', 'requests>=2.0'). Only include packages not in the standard library."
        }
    },
    "required": ["code", "code_type", "docstring", "dependencies", "packages"]
}


# =============================================================================
# User-Friendly Exception Classes
# =============================================================================

class ScryRunError(Exception):
    """Base exception for scry-run errors with user-friendly messages."""
    
    def __init__(self, message: str, hint: Optional[str] = None):
        self.message = message
        self.hint = hint
        super().__init__(self._format_message())
    
    def _format_message(self) -> str:
        msg = f"[scry-run] {self.message}"
        if self.hint:
            msg += f"\n         💡 {self.hint}"
        return msg


class APIKeyError(ScryRunError):
    """API key is missing or invalid."""
    pass


class RateLimitError(ScryRunError):
    """Rate limit exceeded."""
    pass


class QuotaExceededError(ScryRunError):
    """API quota exhausted."""
    pass


class ModelNotFoundError(ScryRunError):
    """Requested model doesn't exist or isn't accessible."""
    pass


class ContentBlockedError(ScryRunError):
    """Content was blocked by safety filters."""
    pass


class NetworkError(ScryRunError):
    """Network connectivity issue."""
    pass


class CodeGenerationError(ScryRunError):
    """Code generation failed."""
    pass


class CodeValidationError(ScryRunError):
    """Generated code failed to parse."""
    pass


def _handle_api_error(error: Exception, model: str) -> ScryRunError:
    """Convert API exceptions to user-friendly errors.

    Args:
        error: The original exception from the backend
        model: The model name being used (for error messages)

    Returns:
        A user-friendly ScryRunError subclass
    """
    error_str = str(error).lower()
    error_type = type(error).__name__

    # Authentication errors
    if "api_key" in error_str or "authentication" in error_str or "401" in error_str:
        return APIKeyError(
            "Invalid API key",
            "Check your API key configuration"
        )
    
    # Rate limiting
    if "rate" in error_str and "limit" in error_str or "429" in error_str or "resource_exhausted" in error_str:
        return RateLimitError(
            "Rate limit exceeded",
            "Wait a moment and try again. Consider using a slower request rate."
        )
    
    # Quota exceeded
    if "quota" in error_str or "billing" in error_str:
        return QuotaExceededError(
            "API quota exhausted",
            "Check your API quota and billing settings"
        )
    
    # Model not found
    if "model" in error_str and ("not found" in error_str or "not exist" in error_str or "404" in error_str):
        return ModelNotFoundError(
            f"Model '{model}' not found",
            f"Try a different model. Set SCRY_MODEL or backend-specific env var."
        )
    
    # Content blocked by safety filters
    if "safety" in error_str or "blocked" in error_str or "harm" in error_str:
        return ContentBlockedError(
            "Content blocked by safety filters",
            "Try rephrasing your class/method names or docstrings"
        )
    
    # Network errors
    if any(x in error_str for x in ["connection", "timeout", "network", "dns", "ssl"]):
        return NetworkError(
            "Network error connecting to backend",
            "Check your internet connection and firewall settings"
        )

    # Permission/access errors
    if "permission" in error_str or "403" in error_str or "access" in error_str:
        return APIKeyError(
            "API access denied",
            "Check your API key permissions"
        )

    # Server errors
    if "500" in error_str or "503" in error_str or "server" in error_str:
        return NetworkError(
            "Backend server error",
            "The API is temporarily unavailable. Try again in a few seconds."
        )
    
    # Fallback: wrap unknown errors
    return CodeGenerationError(
        f"API error: {error_type}: {error}",
        "Check the error message above. If persistent, file an issue."
    )


class GenerationLimitError(ScryRunError):
    """Global generation limit exceeded."""
    pass


class CodeGenerator:
    """Generates code using pluggable backends.

    The generator maintains a persistent session with the backend.
    Full codebase context is sent once at startup, then each generation
    request sends only the minimal "frame" (class/method info).

    Configure via SCRY_BACKEND env var or backend parameter.
    """

    MAX_RETRIES = 5
    
    # Global counter for all generations in this process
    _total_generations_count = 0
    
    # Track context per-generator instance
    _context_hash: Optional[str] = None
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        model: Optional[str] = None,
        backend: Optional[str] = None,
    ):
        """Initialize the code generator.

        Args:
            api_key: Deprecated, not used.
            model: Model to use. Defaults to backend-specific default.
            backend: Backend to use: "claude", "frozen", or "auto" (default).
                     Can also be set via SCRY_BACKEND env var.
        """
        from scry_run.backends import get_backend

        # Pass model through - backends handle their own defaults
        self.model = model or os.environ.get("SCRY_MODEL")
        self._backend = get_backend(name=backend, model=self.model)
        self._context_provided = False

    
    def _format_installed_packages(self, packages: list[str] | None) -> str:
        """Format installed packages for the prompt."""
        if not packages:
            return "No packages are currently installed in this app's environment."

        pkg_list = "\n".join(f"- {pkg}" for pkg in packages)
        return f"""The following packages are already installed:
{pkg_list}

Prefer using these before requesting new packages. Only add to "packages" if you need something not listed above."""

    def _build_context_prompt(
        self,
        context: str,
        installed_packages: list[str] | None = None,
    ) -> str:
        """Build the full context prompt (sent once at session start).
        
        Args:
            context: Full codebase context with hole marker
            installed_packages: Packages already installed
            
        Returns:
            The formatted context string for the system prompt
        """
        return f"""## Codebase Context

The following is the full codebase you are working with:

```python
{context}
```

## Installed Packages

{self._format_installed_packages(installed_packages)}

## Critical Requirements for Code Generation

### Code Quality
1. The code MUST be valid Python that can be executed with `exec()`
2. Complete method/property definition starting with `def` or `@property`
3. Use appropriate decorators (@classmethod, @staticmethod, @property) if needed
4. STRICTLY USE TYPE HINTS for all parameters and return values

### Application-Level Thinking
5. READ THE CLASS DOCSTRING - it describes what the application does
6. EXTRAPOLATE sensibly based on the application's purpose
7. Generate code that a HUMAN would write for this application

### Dependencies
8. List ALL imports needed in the dependencies field
9. Prefer standard library when possible
"""
    
    def _build_frame_prompt(
        self,
        class_name: str,
        attr_name: str,
        is_classmethod: bool = False,
        runtime_context: str = "",
    ) -> str:
        """Build the minimal frame prompt (sent for each generation).
        
        This assumes the full codebase has already been provided.
        Runtime context (call stack, variables) is included since it's dynamic.
        
        Args:
            class_name: Name of the class needing the attribute
            attr_name: Name of the attribute to generate
            is_classmethod: Whether this is a class-level attribute
            runtime_context: Dynamic context (call stack, local variables)
            
        Returns:
            The minimal prompt for this specific generation
        """
        attr_type = "class method/property" if is_classmethod else "instance method/property"
        
        # Include runtime context if provided
        context_section = ""
        if runtime_context:
            context_section = f"\n## Runtime Context\n\n{runtime_context}\n"
        
        return f"""Generate the `{attr_name}` {attr_type} for the `{class_name}` class.
{context_section}
Respond with a JSON object containing:
- code: The Python code (complete method/property definition with type hints)
- code_type: One of "method", "property", "classmethod", "staticmethod"
- docstring: Brief description of what this code does
- dependencies: List of import statements needed
- packages: List of PyPI packages to install (if any external packages needed)

IMPORTANT: Output ONLY the JSON object, no markdown, no explanation, no code fences."""

    def _build_prompt(
        self,
        context: str,
        class_name: str,
        attr_name: str,
        is_classmethod: bool = False,
        installed_packages: list[str] | None = None,
    ) -> str:
        """Build the full prompt for code generation (legacy/fallback).
        
        This is used when the backend doesn't support set_context().
        
        Args:
            context: Full codebase context with hole marker
            class_name: Name of the class needing the attribute
            attr_name: Name of the attribute to generate
            is_classmethod: Whether this is a class-level attribute
            
        Returns:
            The formatted prompt string
        """
        attr_type = "class method/property" if is_classmethod else "instance method/property"
        
        return f"""You are an expert Python code generator for the scry-run dynamic code generation system.

## Your Role

You are generating code for an APPLICATION, not just isolated functions. Read the class docstring carefully - it describes the PURPOSE of the entire application. Use this to EXTRAPOLATE what this method should do.

## Codebase Context

The following is the full codebase. There is a marker showing where code is missing:

```python
{context}
```

## Task

Generate the `{attr_name}` {attr_type} for the `{class_name}` class.

## Critical Requirements

### Code Quality
1. The code MUST be valid Python that can be executed with `exec()`
2. Complete method/property definition starting with `def` or `@property`
3. Use appropriate decorators (@classmethod, @staticmethod, @property) if needed
4. STRICTLY USE TYPE HINTS for all parameters and return values (e.g. `def foo(x: int) -> str:`)
   - Import `typing` types (List, Dict, Optional, Any) as needed and add to dependencies

### Rich Contextual Comments
4. Include a DETAILED docstring explaining:
   - What this method does
   - How it fits into the overall application described in the class docstring
   - Any assumptions made about the application's behavior
5. Add inline comments for non-obvious logic
6. Document parameter types and return types

### Application-Level Thinking
7. READ THE CLASS DOCSTRING - it describes what the application does
8. EXTRAPOLATE sensibly: if the app is a "todo list", a `save()` method should persist todos
9. Infer reasonable behavior from:
   - The application description (class docstring)
   - The method/property name
   - Other existing methods
   - How similar applications typically work
10. Generate code that a HUMAN would write for this application, not minimal stubs
11. Think about what OTHER PARAMETERS this method might need beyond the obvious ones:
    - Configuration options (e.g., format, encoding, verbosity)
    - Optional filters or modifiers
    - Callbacks or hooks for extensibility
    - Make these optional with sensible defaults

### Dependencies
11. List ALL imports needed in the dependencies field
12. Prefer standard library when possible
13. The code should work standalone once dependencies are imported
14. Prefer standard library when possible
15. The code should work standalone once dependencies are imported

### Packages

The app has its own virtual environment. If you need external packages:
1. Add import statements to "dependencies" (e.g., "import pygame")
2. Add PyPI package names to "packages" (e.g., "pygame")

{self._format_installed_packages(installed_packages)}

## Output Format

Respond with a JSON object containing:
- code: The Python code (complete method/property definition with rich comments)
- code_type: One of "method", "property", "classmethod", "staticmethod"
- docstring: Brief description of what this code does
- dependencies: List of import statements needed
- packages: List of PyPI packages to install (if any external packages needed)

IMPORTANT: Output ONLY the JSON object, no markdown, no explanation, no code fences."""
    
    def _fix_indentation(self, code: str) -> str:
        """Fix common indentation issues in generated code.
        
        The LLM sometimes returns code with leading indentation or
        inconsistent whitespace. This method normalizes it.
        
        Args:
            code: Raw generated code
            
        Returns:
            Code with fixed indentation
        """
        import textwrap
        from scry_run.logging import get_logger
        logger = get_logger()
        
        original_code = code
        
        # Strip leading/trailing empty lines
        lines = code.split('\n')
        while lines and not lines[0].strip():
            lines.pop(0)
        while lines and not lines[-1].strip():
            lines.pop()
        
        if not lines:
            return code
        
        # Check if first non-empty line is indented
        first_line = lines[0]
        if first_line and first_line[0] in ' \t':
            # Use textwrap.dedent to remove common leading whitespace
            code = textwrap.dedent('\n'.join(lines))
            logger.debug("Fixed indentation in generated code", f"ORIGINAL:\n{original_code}\n\nFIXED:\n{code}")
        else:
            code = '\n'.join(lines)
        
        return code
    
    def _validate_code(self, code: str) -> None:
        """Validate that the code parses correctly.
        
        Args:
            code: Python code to validate
            
        Raises:
            CodeValidationError: If code fails to parse
        """
        try:
            ast.parse(code)
        except SyntaxError as e:
            raise CodeValidationError(f"Generated code has syntax error: {e}")

    
    def _parse_response(self, response_text: str) -> GeneratedCode:
        """Parse the JSON response from the LLM.
        
        Args:
            response_text: Raw response text from LLM
            
        Returns:
            GeneratedCode object
            
        Raises:
            CodeGenerationError: If response cannot be parsed
        """
        try:
            data = json.loads(response_text)
        except json.JSONDecodeError as e:
            raise CodeGenerationError(f"Failed to parse JSON response: {e}")
        
        required_fields = ["code", "code_type", "docstring", "dependencies", "packages"]
        for field in required_fields:
            if field not in data:
                raise CodeGenerationError(f"Missing required field: {field}")

        return GeneratedCode(
            code=data["code"],
            code_type=data["code_type"],
            docstring=data["docstring"],
            dependencies=data["dependencies"],
            packages=data["packages"],
        )
    def generate(
        self,
        context: str,
        class_name: str,
        attr_name: str,
        is_classmethod: bool = False,
        installed_packages: list[str] | None = None,
    ) -> GeneratedCode:
        """Generate code for a missing attribute.

        Args:
            context: Full codebase context with hole marker
            class_name: Name of the class needing the attribute
            attr_name: Name of the attribute to generate
            is_classmethod: Whether this is a class-level attribute
            installed_packages: List of packages already installed in the app's venv

        Returns:
            GeneratedCode with the generated code and metadata

        Raises:
            CodeGenerationError: If generation fails after retries
            CodeValidationError: If generated code doesn't parse
            APIKeyError: If API key is invalid
            RateLimitError: If rate limit is exceeded
            QuotaExceededError: If quota is exhausted
            NetworkError: If network issues occur
            GenerationLimitError: If global generation limit is exceeded
        """
        # Check global generation limit
        max_limit_str = os.environ.get("SCRY_MAX_GENERATIONS", "100")
        if max_limit_str:
            try:
                limit = int(max_limit_str)
                if CodeGenerator._total_generations_count >= limit:
                    raise GenerationLimitError(
                        f"Global generation limit of {limit} reached.",
                        f"Increase SCRY_MAX_GENERATIONS env var if this is intended."
                    )
            except ValueError:
                # Ignore invalid integer values in env var
                pass

        # Increment counter
        CodeGenerator._total_generations_count += 1

        # Log generation start
        from scry_run.logging import get_logger
        logger = get_logger()
        
        # Set generation context for frozen backend error messages
        if hasattr(self._backend, 'set_generation_context'):
            self._backend.set_generation_context(class_name, attr_name)

        # Check if backend supports persistent context
        supports_context = hasattr(self._backend, 'set_context')
        
        # Extract runtime context from the full context (it's appended at the end)
        # The runtime context is dynamic and must be included in every request
        runtime_context = ""
        if "# === RUNTIME CONTEXT ===" in context:
            idx = context.find("# === RUNTIME CONTEXT ===")
            runtime_context = context[idx:]
        
        if supports_context and not self._context_provided:
            # First generation: send full context to backend (without runtime context)
            # Runtime context is dynamic and will be sent with each frame
            static_context = context[:context.find("# === RUNTIME CONTEXT ===")] if runtime_context else context
            context_prompt = self._build_context_prompt(static_context, installed_packages)
            self._backend.set_context(context_prompt)
            self._context_provided = True
            logger.info("Generator: Provided full context to backend session")
            
            # For this generation, send frame with runtime context
            prompt = self._build_frame_prompt(class_name, attr_name, is_classmethod, runtime_context)
        elif supports_context:
            # Subsequent generations: send frame with runtime context
            prompt = self._build_frame_prompt(class_name, attr_name, is_classmethod, runtime_context)
        else:
            # Backend doesn't support persistent context, use full prompt
            prompt = self._build_prompt(context, class_name, attr_name, is_classmethod, installed_packages)
        
        logger.log_generation_start(class_name, attr_name, prompt)

        # Delegate to backend with retries
        last_error: Optional[Exception] = None
        previous_errors: list[str] = []  # Track errors for feedback

        for attempt in range(self.MAX_RETRIES):
            # Build error feedback for retry attempts
            error_feedback = ""
            if previous_errors:
                error_feedback = f"""

## PREVIOUS GENERATION FAILED

This is generation attempt #{attempt + 1}. Previous attempt(s) caused errors:

{chr(10).join(f"- Attempt #{i+1}: {err}" for i, err in enumerate(previous_errors))}

Please analyze these errors and generate FIXED code that avoids them.
"""

            # Rebuild prompt with error feedback if needed
            if error_feedback:
                if supports_context:
                    prompt = self._build_frame_prompt(class_name, attr_name, is_classmethod, runtime_context + error_feedback)
                else:
                    prompt = self._build_prompt(context + error_feedback, class_name, attr_name, is_classmethod, installed_packages)
                logger.log_generation_start(class_name, attr_name, prompt)

            try:
                result = self._backend.generate_code(prompt)
                
                # Log raw response
                logger.log_generation_response(result.code if hasattr(result, 'code') else str(result))
                
                # Convert backend result to our GeneratedCode type
                from scry_run.backends.base import GeneratedCode as BackendResult
                if isinstance(result, BackendResult):
                    result = GeneratedCode(
                        code=result.code,
                        code_type=result.code_type,
                        docstring=result.docstring,
                        dependencies=result.dependencies,
                        packages=result.packages,
                    )

                # Fix indentation issues before validation
                result = GeneratedCode(
                    code=self._fix_indentation(result.code),
                    code_type=result.code_type,
                    docstring=result.docstring,
                    dependencies=result.dependencies,
                    packages=result.packages,
                )
                
                # Validate the generated code
                self._validate_code(result.code)

                
                # Log success
                logger.log_generation_success(class_name, attr_name, result.code)
                
                return result
                
            except (CodeValidationError, CodeGenerationError) as e:
                last_error = e
                # Log the validation error with full details
                code_str = result.code if 'result' in dir() and hasattr(result, 'code') else None
                logger.log_validation_error(e, code_str)

                # Track error for feedback on next attempt
                error_msg = f"{type(e).__name__}: {e}"
                if code_str:
                    # Include a snippet of the problematic code
                    code_preview = code_str[:200] + "..." if len(code_str) > 200 else code_str
                    error_msg += f"\n  Generated code (preview): {code_preview}"
                previous_errors.append(error_msg)

                if attempt >= self.MAX_RETRIES - 1:
                    raise
                warning(f"Validation error, retrying ({attempt + 1}/{self.MAX_RETRIES})...")
            except Exception as e:
                # Log the error
                logger.error(f"Backend error: {e}", str(e))

                # Convert to friendly error
                friendly_error = _handle_api_error(e, self.model)
                last_error = friendly_error

                if isinstance(friendly_error, (RateLimitError, NetworkError)):
                    wait_time = 2 ** attempt
                    warning(f"{friendly_error.message}. Retrying in {wait_time:.1f}s...")
                    time.sleep(wait_time)
                else:
                    raise friendly_error
        
        raise CodeGenerationError(
            f"Failed to generate valid code after {self.MAX_RETRIES} attempts",
            f"Last error: {last_error}"
        )


    def generate_freeform(
        self,
        prompt: str,
        system_instruction: Optional[str] = "You are a helpful coding assistant.",
    ) -> str:
        """Generate freeform text from a prompt.
        
        Args:
            prompt: User prompt
            system_instruction: System instruction to guide the model behavior
            
        Returns:
            Generated text content
            
        Raises:
            ScryRunError: On API errors
        """
        # Prepend system instruction if provided
        full_prompt = prompt
        if system_instruction:
            full_prompt = f"{system_instruction}\n\n{prompt}"
        
        try:
            return self._backend.generate_text(full_prompt)
        except Exception as e:
            raise _handle_api_error(e, self.model)
