"""Claude SDK session service for Cleave TUI.

Wraps ClaudeSDKClient for async message streaming with the Claude API.
Handles authentication, session management, and skill discovery.
"""

from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any

from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient, query
from claude_agent_sdk.types import AssistantMessage, ResultMessage, SystemMessage


class SessionState(Enum):
    """SDK session state."""

    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    ERROR = "error"


@dataclass
class Message:
    """A message from the SDK session."""

    role: str  # "user", "assistant", "system"
    content: str
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class SDKConfig:
    """Configuration for the SDK session."""

    model: str = "claude-sonnet-4-20250514"
    max_turns: int = 50
    working_directory: Path | None = None
    system_prompt: str | None = None
    allowed_tools: list[str] | None = None
    disallowed_tools: list[str] | None = None


# System prompts for Calf workflows - brevity is paramount

RAISE_CALF_PROMPT = """\
Help define a directive for Cleave (task decomposition). Be terse.

Ask 1-2 questions max per turn. Focus on:
- Goal (what)
- Success criteria (how we know it worked)
- Constraints (tech stack, compatibility)
- Out of scope (what NOT to do)

User speaks first. Narrow them down efficiently. No fluff.

When ready: show directive, get confirmation. On confirm, output only:
DIRECTIVE: <the directive text>

Occasional levity permitted. The pigs running this operation are... industrious.
"""

FATTEN_CALF_PROMPT = """\
Refine an existing Cleave directive. The current version:

---
{directive}
---

Help sharpen it. Ask 1-2 questions max per turn:
- Ambiguities to resolve?
- Missing constraints?
- Scope creep to trim?
- Success criteria unclear?

User speaks first. Be surgical. No padding.

When ready: show revised directive, get confirmation. On confirm, output only:
DIRECTIVE: <the refined directive text>

The swine have reviewed the docket. Efficiency is... expected.
"""


class SDKSession:
    """Claude SDK session wrapper.

    Provides async streaming for SDK interactions, supporting both
    interrogation (Calf Raises) and execution (Cleave) modes.

    Usage:
        session = SDKSession(config)

        # Start session
        await session.connect()

        # Send messages and stream responses
        async for message in session.query("What should I build?"):
            print(message.content)

        # Close when done
        await session.disconnect()
    """

    def __init__(self, config: SDKConfig | None = None):
        """Initialize SDK session.

        Args:
            config: Session configuration. Uses defaults if not provided.
        """
        self.config = config or SDKConfig()
        self._client: ClaudeSDKClient | None = None
        self._state = SessionState.DISCONNECTED
        self._error: str | None = None

    @property
    def state(self) -> SessionState:
        """Current session state."""
        return self._state

    @property
    def error(self) -> str | None:
        """Last error message, if any."""
        return self._error

    def is_authenticated(self) -> bool:
        """Check if SDK is authenticated.

        The SDK uses Claude CLI authentication stored in ~/.claude/.
        If not authenticated, SDK operations will fail.

        Returns:
            True if authentication credentials exist.
        """
        # Check for Claude Code auth file
        auth_file = Path.home() / ".claude" / ".credentials.json"
        return auth_file.exists()

    @staticmethod
    def check_auth() -> tuple[bool, str]:
        """Check authentication status with detailed message.

        Returns:
            Tuple of (is_authenticated, message).
        """
        auth_file = Path.home() / ".claude" / ".credentials.json"
        if auth_file.exists():
            return True, "Authenticated"
        return False, "Not authenticated. Run 'claude' CLI to login via browser."

    async def authenticate(self) -> bool:
        """Trigger authentication flow if needed.

        Returns:
            True if authentication succeeded.
        """
        # The SDK handles auth automatically via Claude CLI
        # This is a placeholder for explicit auth handling if needed
        return self.is_authenticated()

    async def connect(self) -> None:
        """Initialize the SDK client connection."""
        if self._state == SessionState.CONNECTED:
            return

        self._state = SessionState.CONNECTING
        self._error = None

        try:
            options = ClaudeAgentOptions(
                max_turns=self.config.max_turns,
            )

            if self.config.system_prompt:
                options.system_prompt = self.config.system_prompt

            if self.config.allowed_tools:
                options.allowed_tools = self.config.allowed_tools

            if self.config.disallowed_tools:
                options.disallowed_tools = self.config.disallowed_tools

            self._client = ClaudeSDKClient(options=options)
            await self._client.__aenter__()
            self._state = SessionState.CONNECTED

        except Exception as e:
            self._state = SessionState.ERROR
            self._error = str(e)
            raise

    async def disconnect(self) -> None:
        """Close the SDK client connection."""
        if self._client is not None:
            try:
                await self._client.__aexit__(None, None, None)
            except Exception:
                pass  # Best effort cleanup
            finally:
                self._client = None
                self._state = SessionState.DISCONNECTED

    async def query(self, prompt: str) -> AsyncIterator[Message]:
        """Send a message and stream responses.

        Args:
            prompt: User message to send.

        Yields:
            Message objects as they arrive from the SDK.

        Raises:
            RuntimeError: If not connected.
        """
        if self._state != SessionState.CONNECTED or self._client is None:
            raise RuntimeError("SDK session not connected. Call connect() first.")

        try:
            await self._client.query(prompt)

            async for event in self._client.receive_response():
                # Parse SDK events into our Message format
                message = self._parse_event(event)
                if message:
                    yield message

        except Exception as e:
            self._state = SessionState.ERROR
            self._error = str(e)
            raise

    def _parse_event(self, event: Any) -> Message | None:
        """Parse SDK event into a Message.

        Args:
            event: Raw SDK event.

        Returns:
            Message if the event contains displayable content, None otherwise.
        """
        # Handle typed SDK events
        if isinstance(event, AssistantMessage):
            # Extract text from content blocks
            text_parts = []
            for block in event.content:
                if hasattr(block, "text"):
                    text_parts.append(block.text)
            if text_parts:
                return Message(role="assistant", content="".join(text_parts))
            return None

        if isinstance(event, SystemMessage):
            # System messages (init, etc.) - usually not displayed
            if event.subtype == "init":
                return None  # Skip init messages
            return Message(
                role="system",
                content=f"[{event.subtype}]",
                metadata=event.data if hasattr(event, "data") else {},
            )

        if isinstance(event, ResultMessage):
            # Final result - could show summary
            if event.is_error:
                return Message(role="system", content=f"[Error: {event.result}]")
            # Don't duplicate the result since AssistantMessage already captured it
            return None

        # Handle legacy/generic event types
        if hasattr(event, "type"):
            if event.type == "text":
                return Message(role="assistant", content=event.text)
            elif event.type == "tool_use":
                return Message(
                    role="system",
                    content=f"[Tool: {event.name}]",
                    metadata={"tool": event.name, "input": getattr(event, "input", {})},
                )

        # Fallback for string events
        if isinstance(event, str):
            return Message(role="assistant", content=event)

        return None


async def simple_query(prompt: str, system_prompt: str | None = None) -> AsyncIterator[Message]:
    """Simple one-shot query without session management.

    Useful for quick queries that don't need multi-turn conversation.

    Args:
        prompt: The query prompt.
        system_prompt: Optional system prompt.

    Yields:
        Message objects from the response.
    """
    try:
        async for event in query(prompt=prompt, system_prompt=system_prompt):
            if isinstance(event, str):
                yield Message(role="assistant", content=event)
            elif hasattr(event, "text"):
                yield Message(role="assistant", content=event.text)
    except Exception as e:
        yield Message(role="system", content=f"Error: {e}")


class InterrogationSession:
    """Specialized session for Calf workflows.

    Two modes:
    - Raise: Create new directive from scratch
    - Fatten: Refine existing directive
    """

    def __init__(
        self,
        working_directory: Path | None = None,
        mode: str = "raise",
        existing_directive: str | None = None,
    ):
        """Initialize interrogation session.

        Args:
            working_directory: Directory for .calf file output.
            mode: "raise" for new directive, "fatten" for editing existing.
            existing_directive: Content of existing .calf file (for fatten mode).
        """
        self.working_directory = working_directory or Path.cwd()
        self.mode = mode
        self.existing_directive = existing_directive
        self._session: SDKSession | None = None
        self._messages: list[Message] = []
        self._directive: str | None = None

    @property
    def is_complete(self) -> bool:
        """Whether interrogation has produced a directive."""
        return self._directive is not None

    @property
    def directive(self) -> str | None:
        """The extracted directive, if complete."""
        return self._directive

    async def start(self) -> AsyncIterator[Message]:
        """Start the interrogation session.

        Yields:
            Opening messages from the assistant.
        """
        # Select prompt based on mode
        if self.mode == "fatten" and self.existing_directive:
            system_prompt = FATTEN_CALF_PROMPT.format(directive=self.existing_directive)
            opening = "What needs tightening?"
        else:
            system_prompt = RAISE_CALF_PROMPT
            opening = "What are we building?"

        config = SDKConfig(
            system_prompt=system_prompt,
            working_directory=self.working_directory,
        )
        self._session = SDKSession(config)
        await self._session.connect()

        # Initial prompt - user speaks first per the system prompt
        async for msg in self._session.query(opening):
            self._messages.append(msg)
            yield msg

    async def respond(self, user_input: str) -> AsyncIterator[Message]:
        """Send user response and get assistant reply.

        Args:
            user_input: User's response to clarifying questions.

        Yields:
            Assistant messages in response.
        """
        if self._session is None:
            raise RuntimeError("Session not started. Call start() first.")

        # Track user message
        self._messages.append(Message(role="user", content=user_input))

        async for msg in self._session.query(user_input):
            self._messages.append(msg)

            # Check if this is the final directive
            if "DIRECTIVE:" in msg.content:
                # Extract everything after DIRECTIVE:
                idx = msg.content.index("DIRECTIVE:")
                self._directive = msg.content[idx + 10:].strip()

            yield msg

    async def close(self) -> None:
        """Close the interrogation session."""
        if self._session:
            await self._session.disconnect()
            self._session = None


class CleaveSkillError(Exception):
    """Raised when the Cleave skill is not available."""

    pass


class CleaveExecutor:
    """Executes directives through the Cleave skill.

    Wraps SDK session with Cleave-specific configuration.
    """

    # Path to user-level Cleave skill
    CLEAVE_SKILL_PATH = Path.home() / ".claude" / "skills" / "cleave"

    def __init__(
        self,
        directive: str,
        model: str = "opus",
        max_depth: int = 3,
        parallel: bool = True,
        dry_run: bool = False,
        working_directory: Path | None = None,
    ):
        """Initialize Cleave executor.

        Args:
            directive: The directive text to execute.
            model: Model to use (sonnet, opus, haiku).
            max_depth: Maximum recursion depth for task decomposition.
            parallel: Whether to enable parallel execution.
            dry_run: If True, assess but don't execute.
            working_directory: Working directory for execution.
        """
        self.directive = directive
        self.model = model
        self.max_depth = max_depth
        self.parallel = parallel
        self.dry_run = dry_run
        self.working_directory = working_directory or Path.cwd()
        self._session: SDKSession | None = None
        self._progress: tuple[int, int] = (0, 0)  # (current, total)
        self._current_task: str = ""

    @classmethod
    def is_cleave_available(cls) -> bool:
        """Check if the Cleave skill is installed.

        Returns:
            True if the Cleave skill directory exists.
        """
        return cls.CLEAVE_SKILL_PATH.exists()

    @classmethod
    def validate_cleave_skill(cls) -> None:
        """Validate that the Cleave skill is available.

        Raises:
            CleaveSkillError: If the skill is not installed.
        """
        if not cls.is_cleave_available():
            raise CleaveSkillError(
                f"Cleave skill not found at {cls.CLEAVE_SKILL_PATH}. "
                "Install the Cleave skill to ~/.claude/skills/cleave/"
            )

    @property
    def progress(self) -> tuple[int, int]:
        """Current progress (current_step, total_steps)."""
        return self._progress

    @property
    def current_task(self) -> str:
        """Description of the current task being executed."""
        return self._current_task

    async def execute(self) -> AsyncIterator[Message]:
        """Execute the directive through Cleave.

        Yields:
            Messages from the execution, including progress updates.

        Raises:
            CleaveSkillError: If the Cleave skill is not installed.
        """
        # Validate Cleave skill is available
        self.validate_cleave_skill()

        # Build the Cleave invocation prompt
        cleave_prompt = self._build_cleave_prompt()

        # Map model names
        model_map = {
            "sonnet": "claude-sonnet-4-20250514",
            "opus": "claude-opus-4-20250514",
            "haiku": "claude-haiku-4-20250514",
        }

        config = SDKConfig(
            model=model_map.get(self.model, model_map["opus"]),
            working_directory=self.working_directory,
        )

        self._session = SDKSession(config)

        try:
            await self._session.connect()

            async for msg in self._session.query(cleave_prompt):
                # Parse progress from output if present
                self._parse_progress(msg.content)
                yield msg

        finally:
            if self._session:
                await self._session.disconnect()

    def _build_cleave_prompt(self) -> str:
        """Build the prompt to invoke Cleave skill.

        Returns:
            Formatted prompt for Cleave execution.
        """
        flags = []
        if self.parallel:
            flags.append("parallel=true")
        if self.dry_run:
            flags.append("dry_run=true")
        flags.append(f"max_depth={self.max_depth}")

        flags_str = ", ".join(flags)

        return f"""/cleave {flags_str}

{self.directive}
"""

    def _parse_progress(self, content: str) -> None:
        """Parse progress information from message content.

        Args:
            content: Message content that may contain progress info.
        """
        # Look for patterns like "3/12 subtasks" or "Step 3 of 12"
        import re

        # Pattern: X/Y format
        match = re.search(r"(\d+)/(\d+)", content)
        if match:
            self._progress = (int(match.group(1)), int(match.group(2)))

        # Pattern: "Current: <task description>"
        match = re.search(r"Current:\s*(.+?)(?:\n|$)", content)
        if match:
            self._current_task = match.group(1).strip()
