"""Claude Code backend - uses persistent ClaudeSDKClient session.

This backend maintains a single Claude process across multiple code generation
requests. Context is provided only on session startup; subsequent generations
send only the minimal "frame" (class/method to generate).
"""

from __future__ import annotations

import asyncio
import atexit
import json
import shutil
import threading
from typing import Any, Optional

from scry_run.backends.base import GeneratorBackend, GeneratedCode
from scry_run.backends.registry import register_backend
from scry_run.console import err_console


class ClaudeSessionManager:
    """Manages a persistent ClaudeSDKClient session.
    
    This singleton maintains a Claude session across multiple requests.
    Context is provided once at startup, then each generation request
    sends only the minimal "frame" information.
    
    Session resilience:
    - If session dies, it will be restarted on next request
    - Full context is re-provided on restart
    """
    
    _instance: Optional["ClaudeSessionManager"] = None
    _lock = threading.Lock()
    
    def __new__(cls) -> "ClaudeSessionManager":
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
        
        self._initialized = True
        self._client: Any = None
        self._connected = False
        self._context: Optional[str] = None  # Stored for session restart
        self._model: Optional[str] = None
        
        # Background event loop for async SDK
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._thread: Optional[threading.Thread] = None
        self._thread_ready = threading.Event()
        
        # Register shutdown handler
        atexit.register(self._shutdown_sync)
    
    def _ensure_loop_running(self) -> None:
        """Ensure the background event loop is running."""
        if self._loop is not None and self._thread is not None and self._thread.is_alive():
            return
        
        def run_loop():
            self._loop = asyncio.new_event_loop()
            asyncio.set_event_loop(self._loop)
            self._thread_ready.set()
            self._loop.run_forever()
        
        self._thread = threading.Thread(target=run_loop, daemon=True, name="claude-sdk-loop")
        self._thread.start()
        self._thread_ready.wait(timeout=5.0)
        
        if self._loop is None:
            raise RuntimeError("Failed to start background event loop")
    
    def _run_async(self, coro, timeout: int = 300) -> Any:
        """Run an async coroutine synchronously using the background loop.
        
        Args:
            coro: Coroutine to run
            timeout: Timeout in seconds (default: 5 minutes)
        """
        self._ensure_loop_running()
        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
        return future.result(timeout=timeout)
    
    def connect(self, context: str, model: Optional[str] = None) -> None:
        """Connect to Claude with the full codebase context.
        
        Args:
            context: Full codebase context (sent once as system prompt)
            model: Model to use (e.g., 'claude-sonnet-4-5')
        """
        from scry_run.logging import get_logger
        logger = get_logger()
        
        # Store for session restart
        self._context = context
        self._model = model
        
        if self._connected:
            return
        
        logger.info("Claude backend: Starting persistent session")
        
        try:
            self._run_async(self._connect_async(context, model))
            self._connected = True
            logger.info("Claude backend: Session connected")
        except Exception as e:
            logger.error(f"Claude backend: Connection failed: {e}")
            raise
    
    async def _connect_async(self, context: str, model: Optional[str]) -> None:
        """Async implementation of connect."""
        from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient
        
        # Provide full context as system prompt - this is sent only once
        system_prompt = 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. The codebase context below describes the application you're working with.
For each code generation request, you will receive:
1. The class name and attribute name to generate
2. Whether it's an instance or class method/property

Use the codebase context to understand the application and generate appropriate code.

## Codebase Context

{context}

## Output Format

For each generation request, respond ONLY 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
- dependencies: List of import statements needed
- packages: List of PyPI packages to install (if any)

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

        options = ClaudeAgentOptions(
            system_prompt=system_prompt,
            model=model,
        )
        
        self._client = ClaudeSDKClient(options=options)
        await self._client.connect()
    
    def _ensure_connected(self) -> None:
        """Ensure session is connected, restart if needed."""
        if self._connected and self._client:
            return
        
        if self._context is None:
            raise RuntimeError("No context provided. Call connect() first.")
        
        # Attempt reconnect with stored context
        from scry_run.logging import get_logger
        logger = get_logger()
        logger.info("Claude backend: Session died, reconnecting with full context...")
        
        self._connected = False
        self.connect(self._context, self._model)
    
    def query(self, prompt: str, max_retries: int = 3) -> str:
        """Send a query and collect the full text response.
        
        Args:
            prompt: The minimal prompt (just the "frame" - class/method info)
            max_retries: Number of retries on session failure
            
        Returns:
            The collected text response from Claude
        """
        from scry_run.logging import get_logger
        logger = get_logger()
        
        last_error: Optional[Exception] = None
        
        for attempt in range(max_retries):
            try:
                self._ensure_connected()
                
                # Truncate prompt for logging
                if len(prompt) > 500:
                    prompt_preview = prompt[:200] + f"\n... [{len(prompt) - 400} chars] ...\n" + prompt[-200:]
                else:
                    prompt_preview = prompt
                logger.debug(f"Claude backend: query ({len(prompt)} chars)", f"PROMPT:\n{prompt_preview}")
                
                response = self._run_async(self._query_async(prompt))
                
                logger.debug(f"Claude backend: response ({len(response)} chars)", f"RESPONSE:\n{response}")
                return response
                
            except Exception as e:
                last_error = e
                logger.error(f"Claude backend: Query failed (attempt {attempt + 1}/{max_retries}): {e}")
                
                # Mark as disconnected to force reconnect on next attempt
                self._connected = False
                self._client = None
                
                if attempt < max_retries - 1:
                    logger.info("Claude backend: Retrying with fresh session...")
        
        raise RuntimeError(f"Claude query failed after {max_retries} attempts: {last_error}")
    
    async def _query_async(self, prompt: str) -> str:
        """Async implementation of query."""
        from claude_agent_sdk import AssistantMessage, ResultMessage, TextBlock
        
        await self._client.query(prompt)
        
        # Collect all text from the response
        text_parts = []
        
        async for message in self._client.receive_response():
            if isinstance(message, AssistantMessage):
                for block in message.content:
                    if isinstance(block, TextBlock):
                        text_parts.append(block.text)
            elif isinstance(message, ResultMessage):
                if message.is_error:
                    raise RuntimeError(f"Claude query failed: {message.result}")
                break
        
        return "".join(text_parts)
    
    def disconnect(self) -> None:
        """Disconnect the session."""
        if not self._connected:
            return
        
        from scry_run.logging import get_logger
        logger = get_logger()
        logger.info("Claude backend: Disconnecting session")
        
        try:
            self._run_async(self._disconnect_async(), timeout=10)
        except Exception as e:
            # Cancel scope errors are expected during atexit (different thread)
            # Just log at debug level and proceed with cleanup
            if "cancel scope" in str(e).lower():
                logger.debug(f"Claude backend: Expected cleanup error: {e}")
            else:
                logger.error(f"Claude backend: Error during disconnect: {e}")
        
        self._connected = False
    
    async def _disconnect_async(self) -> None:
        """Async implementation of disconnect."""
        if self._client:
            try:
                await self._client.disconnect()
            except Exception:
                # Ignore errors during disconnect - we're cleaning up anyway
                pass
            self._client = None
    
    def reset(self) -> None:
        """Reset the session (forces full context on next connect)."""
        self.disconnect()
        self._context = None
    
    def _shutdown_sync(self) -> None:
        """Synchronous shutdown for atexit handler."""
        # Mark as disconnected first to prevent further operations
        self._connected = False
        
        # Stop the event loop - this is the cleanest way to shut down
        # Don't try to run async disconnect as it causes task issues
        if self._loop is not None:
            try:
                self._loop.call_soon_threadsafe(self._loop.stop)
            except Exception:
                pass  # Loop may already be stopped
    
    @property
    def is_connected(self) -> bool:
        """Check if session is connected."""
        return self._connected
    
    @property
    def has_context(self) -> bool:
        """Check if context has been provided."""
        return self._context is not None


class ClaudeBackend(GeneratorBackend):
    """Backend that uses persistent ClaudeSDKClient session.
    
    Context is provided once at session startup. Subsequent generation
    requests send only the minimal "frame" information.
    """
    
    name = "claude"
    
    def __init__(self, model: Optional[str] = None):
        """Initialize Claude backend.
        
        Args:
            model: Model name (e.g., 'claude-sonnet-4-5')
        """
        self.model = model
        self._session = ClaudeSessionManager()
    
    @classmethod
    def is_available(cls) -> bool:
        """Check if claude-agent-sdk is available."""
        try:
            import claude_agent_sdk
            return True
        except ImportError:
            return shutil.which("claude") is not None
    
    def set_context(self, context: str) -> None:
        """Set the codebase context for this session.
        
        This should be called once before any generation requests.
        The context will be sent as the system prompt.
        
        Args:
            context: Full codebase context
        """
        self._session.connect(context, self.model)
    
    def generate_text(self, prompt: str) -> str:
        """Generate freeform text via persistent session.
        
        Note: For best results, call set_context() first to establish
        the codebase context.
        """
        from scry_run.logging import get_logger
        logger = get_logger()
        logger.info("Claude backend: generate_text called")
        
        # If no context yet, just send the prompt directly
        if not self._session.has_context:
            self._session.connect("", self.model)
        
        return self._session.query(prompt)
    
    def generate_code(self, prompt: str) -> GeneratedCode:
        """Generate structured code via persistent session.
        
        The prompt should be the minimal "frame" - which class/method
        to generate. The full context should have been set via set_context().
        
        If set_context() wasn't called, the prompt will be used as-is
        (for backward compatibility).
        """
        from scry_run.logging import get_logger
        logger = get_logger()
        logger.info("Claude backend: generate_code called")
        
        # If no context set, use prompt as context (backward compat)
        if not self._session.has_context:
            # Extract context from prompt and use it
            self._session.connect(prompt, self.model)
            # First call sets context, so we send a minimal follow-up
            generation_prompt = "Generate the code as specified in the context above."
        else:
            # Context already set, send just the generation request
            generation_prompt = prompt
        
        response = self._session.query(generation_prompt)
        
        # Parse JSON response
        try:
            start = response.find("{")
            end = response.rfind("}") + 1
            if start >= 0 and end > start:
                json_str = response[start:end]
                logger.debug(f"Claude backend: extracted JSON ({len(json_str)} chars)", f"JSON:\n{json_str}")
                data = json.loads(json_str)
            else:
                data = json.loads(response)
        except json.JSONDecodeError as e:
            logger.error(f"Claude backend: JSON parse error", f"ERROR: {e}\n\nRESPONSE:\n{response}")
            raise RuntimeError(f"Failed to parse JSON from response: {e}\nResponse: {response[:500]}")
        
        logger.info(f"Claude backend: code generated (type={data.get('code_type', 'unknown')}, {len(data.get('code', ''))} chars)")
        
        return GeneratedCode(
            code=data.get("code", ""),
            code_type=data.get("code_type", "method"),
            docstring=data.get("docstring", ""),
            dependencies=data.get("dependencies", []),
            packages=data.get("packages", []),
        )


# Register on import
register_backend("claude", ClaudeBackend)
