"""
Session management with markdown file persistence.

Sessions are stored as markdown files in the sessions/ directory.
Each file contains the full conversation history in a readable format.
"""

import os
import re
import json
from pathlib import Path
from datetime import datetime
from typing import Optional, List, Dict, Any
import uuid


# Sessions directory
SESSIONS_DIR = Path(__file__).parent.parent / "sessions"
SESSIONS_DIR.mkdir(exist_ok=True)


class Session:
    """A conversation session persisted to markdown."""

    def __init__(self, session_id: str = None, title: str = None):
        self.session_id = session_id or datetime.now().strftime("%Y%m%d-%H%M%S") + "-" + str(uuid.uuid4())[:8]
        self.title = title or "Untitled Session"
        self.messages: List[Dict[str, Any]] = []
        self.created_at = datetime.utcnow().isoformat()
        self.last_active = self.created_at
        self.file_path = SESSIONS_DIR / f"{self.session_id}.md"

    def add_user_message(self, content: str):
        """Add a user message and save to file."""
        self.messages.append({
            "role": "user",
            "content": content,
            "timestamp": datetime.utcnow().isoformat()
        })
        self.last_active = datetime.utcnow().isoformat()

        # Update title from first message if untitled
        if self.title == "Untitled Session" and len(self.messages) == 1:
            self.title = content[:50] + ("..." if len(content) > 50 else "")

        self._save()

    def add_assistant_message(self, content: Any):
        """Add an assistant message and save to file."""
        self.messages.append({
            "role": "assistant",
            "content": content,
            "timestamp": datetime.utcnow().isoformat()
        })
        self.last_active = datetime.utcnow().isoformat()
        self._save()

    def add_tool_results(self, results: List[Dict]):
        """Add tool results and save to file."""
        self.messages.append({
            "role": "user",
            "content": results,
            "timestamp": datetime.utcnow().isoformat()
        })
        self.last_active = datetime.utcnow().isoformat()
        self._save()

    def get_messages_for_api(self) -> List[Dict[str, Any]]:
        """Get messages in the format expected by the Anthropic API."""
        api_messages = []
        for msg in self.messages:
            api_messages.append({
                "role": msg["role"],
                "content": msg["content"]
            })
        return api_messages

    def _save(self):
        """Save session to markdown file."""
        md_content = self._to_markdown()
        self.file_path.write_text(md_content)

    def _to_markdown(self) -> str:
        """Convert session to markdown format."""
        lines = []

        # Frontmatter
        lines.append("---")
        lines.append(f"session_id: {self.session_id}")
        lines.append(f"title: \"{self.title}\"")
        lines.append(f"created_at: {self.created_at}")
        lines.append(f"last_active: {self.last_active}")
        lines.append(f"message_count: {len(self.messages)}")
        lines.append("---")
        lines.append("")
        lines.append(f"# {self.title}")
        lines.append("")

        # Messages
        for i, msg in enumerate(self.messages):
            role = msg["role"]
            content = msg["content"]
            timestamp = msg.get("timestamp", "")

            if role == "user":
                if isinstance(content, str):
                    lines.append(f"## User")
                    lines.append(f"*{timestamp}*")
                    lines.append("")
                    lines.append(content)
                    lines.append("")
                elif isinstance(content, list):
                    # Tool results
                    lines.append(f"## Tool Results")
                    lines.append(f"*{timestamp}*")
                    lines.append("")
                    for result in content:
                        if isinstance(result, dict) and result.get("type") == "tool_result":
                            tool_id = result.get("tool_use_id", "unknown")
                            result_content = result.get("content", "")
                            lines.append(f"<details>")
                            lines.append(f"<summary>Result for {tool_id[:12]}...</summary>")
                            lines.append("")
                            lines.append("```")
                            lines.append(result_content[:2000])
                            if len(result_content) > 2000:
                                lines.append("... (truncated)")
                            lines.append("```")
                            lines.append("</details>")
                            lines.append("")

            elif role == "assistant":
                lines.append(f"## Assistant")
                lines.append(f"*{timestamp}*")
                lines.append("")

                if isinstance(content, str):
                    lines.append(content)
                elif isinstance(content, list):
                    # Handle content blocks (text, tool_use)
                    for block in content:
                        if hasattr(block, "type"):
                            if block.type == "text" and block.text:
                                lines.append(block.text)
                                lines.append("")
                            elif block.type == "tool_use":
                                lines.append(f"### Tool Call: `{block.name}`")
                                lines.append("")
                                lines.append("```json")
                                lines.append(json.dumps(block.input, indent=2))
                                lines.append("```")
                                lines.append("")
                        elif isinstance(block, dict):
                            if block.get("type") == "text":
                                lines.append(block.get("text", ""))
                                lines.append("")
                            elif block.get("type") == "tool_use":
                                lines.append(f"### Tool Call: `{block.get('name')}`")
                                lines.append("")
                                lines.append("```json")
                                lines.append(json.dumps(block.get("input", {}), indent=2))
                                lines.append("```")
                                lines.append("")
                lines.append("")

        return "\n".join(lines)

    @classmethod
    def load(cls, session_id: str) -> Optional["Session"]:
        """Load a session from markdown file."""
        file_path = SESSIONS_DIR / f"{session_id}.md"
        if not file_path.exists():
            return None

        content = file_path.read_text()
        return cls._from_markdown(content, session_id)

    @classmethod
    def _from_markdown(cls, content: str, session_id: str) -> "Session":
        """Parse markdown back into a session (basic reconstruction)."""
        session = cls(session_id=session_id)

        # Parse frontmatter
        if content.startswith("---"):
            parts = content.split("---", 2)
            if len(parts) >= 3:
                frontmatter = parts[1]
                for line in frontmatter.strip().split("\n"):
                    if line.startswith("title:"):
                        session.title = line.split(":", 1)[1].strip().strip('"')
                    elif line.startswith("created_at:"):
                        session.created_at = line.split(":", 1)[1].strip()
                    elif line.startswith("last_active:"):
                        session.last_active = line.split(":", 1)[1].strip()

        # Note: Full message reconstruction from markdown is complex
        # For now, we store the raw JSON alongside for accurate reconstruction
        json_path = SESSIONS_DIR / f"{session_id}.json"
        if json_path.exists():
            try:
                data = json.loads(json_path.read_text())
                session.messages = data.get("messages", [])
            except:
                pass

        session.file_path = SESSIONS_DIR / f"{session_id}.md"
        return session

    def _save(self):
        """Save session to both markdown and JSON files."""
        # Save markdown for readability
        md_content = self._to_markdown()
        self.file_path.write_text(md_content)

        # Save JSON for accurate reconstruction
        json_path = SESSIONS_DIR / f"{self.session_id}.json"
        json_data = {
            "session_id": self.session_id,
            "title": self.title,
            "created_at": self.created_at,
            "last_active": self.last_active,
            "messages": self._serialize_messages()
        }
        json_path.write_text(json.dumps(json_data, indent=2, default=str))

    def _serialize_messages(self) -> List[Dict]:
        """Serialize messages for JSON storage."""
        serialized = []
        for msg in self.messages:
            content = msg["content"]

            # Handle Anthropic content blocks
            if isinstance(content, list):
                serialized_content = []
                for block in content:
                    if hasattr(block, "type"):
                        # Anthropic SDK object
                        if block.type == "text":
                            serialized_content.append({"type": "text", "text": block.text})
                        elif block.type == "tool_use":
                            serialized_content.append({
                                "type": "tool_use",
                                "id": block.id,
                                "name": block.name,
                                "input": block.input
                            })
                    else:
                        serialized_content.append(block)
                content = serialized_content

            serialized.append({
                "role": msg["role"],
                "content": content,
                "timestamp": msg.get("timestamp", "")
            })

        return serialized


# Session store (in-memory cache backed by files)
_sessions: Dict[str, Session] = {}


def get_or_create_session(session_id: Optional[str] = None) -> Session:
    """Get existing session or create new one."""
    if session_id:
        # Check memory cache first
        if session_id in _sessions:
            return _sessions[session_id]

        # Try to load from file
        session = Session.load(session_id)
        if session:
            _sessions[session_id] = session
            return session

    # Create new session
    session = Session(session_id)
    _sessions[session.session_id] = session
    return session


def get_session(session_id: str) -> Optional[Session]:
    """Get session by ID."""
    if session_id in _sessions:
        return _sessions[session_id]

    session = Session.load(session_id)
    if session:
        _sessions[session_id] = session
    return session


def list_sessions() -> List[Session]:
    """List all sessions from files."""
    sessions = []

    for md_file in SESSIONS_DIR.glob("*.md"):
        session_id = md_file.stem
        session = get_session(session_id)
        if session:
            sessions.append(session)

    # Sort by last_active descending
    sessions.sort(key=lambda s: s.last_active, reverse=True)
    return sessions


def delete_session(session_id: str) -> bool:
    """Delete a session."""
    if session_id in _sessions:
        del _sessions[session_id]

    md_path = SESSIONS_DIR / f"{session_id}.md"
    json_path = SESSIONS_DIR / f"{session_id}.json"

    deleted = False
    if md_path.exists():
        md_path.unlink()
        deleted = True
    if json_path.exists():
        json_path.unlink()
        deleted = True

    return deleted
