"""
State management for chunk tracking.

Tracks the start position of each task chunk so we know
which portion of the JSONL to triage when the task completes.

NEW ARCHITECTURE (v2):
- State files are now per-project AND per-session
- Location: ~/.codecontext/projects/{encoded-path}/state/chunk_{session_id}.json
- This prevents concurrent session collisions (the critical bug in v1)
"""

import json
import sys
from pathlib import Path
from datetime import datetime
from typing import Optional
from dataclasses import dataclass, asdict

# Add codecontext to path for imports (hooks/ is inside codecontext/)
sys.path.insert(0, str(Path(__file__).parent.parent))

from cli_utils import (
    CODECONTEXT_ROOT,
    PROJECTS_DIR,
    encode_project_path,
    get_state_dir
)


@dataclass
class ChunkState:
    """State for tracking a task chunk."""

    chunk_start_timestamp: str
    session_id: str
    project_path: str
    user_prompt: str
    jsonl_path: Optional[str] = None

    @classmethod
    def create(
        cls,
        session_id: str,
        project_path: str,
        user_prompt: str,
        jsonl_path: Optional[str] = None
    ) -> "ChunkState":
        """Create a new chunk state."""
        return cls(
            # Use UTC to match JSONL timestamps (which use Z suffix = UTC)
            chunk_start_timestamp=datetime.utcnow().isoformat() + "Z",
            session_id=session_id,
            project_path=project_path,
            user_prompt=user_prompt,
            jsonl_path=jsonl_path
        )

    def _get_state_file(self) -> Path:
        """Get the state file path for this chunk's project and session."""
        state_dir = get_state_dir(self.project_path)
        return state_dir / f"chunk_{self.session_id}.json"

    def save(self) -> Path:
        """
        Save state to per-project, per-session file.

        State file location:
            ~/.codecontext/projects/{encoded-path}/state/chunk_{session_id}.json

        This ensures:
            - No collision between different projects
            - No collision between concurrent sessions in same project
        """
        state_file = self._get_state_file()
        state_file.parent.mkdir(parents=True, exist_ok=True)
        state_file.write_text(json.dumps(asdict(self), indent=2))
        return state_file

    @classmethod
    def load(cls, session_id: str = None, project_path: str = None) -> Optional["ChunkState"]:
        """
        Load state from file.

        Args:
            session_id: Session ID to load state for (optional)
            project_path: Project path to narrow search (optional)

        Returns:
            ChunkState if found, None otherwise

        Search strategy:
            1. If both session_id and project_path provided: direct lookup
            2. If only session_id: search all projects
            3. If only project_path: find most recent state in project
            4. If neither: find most recent state globally (legacy behavior)
        """
        # Strategy 1: Direct lookup with both parameters
        if session_id and project_path:
            state_dir = get_state_dir(project_path)
            state_file = state_dir / f"chunk_{session_id}.json"
            if state_file.exists():
                return cls._load_from_file(state_file)

        # Strategy 2: Search all projects for session_id
        if session_id and not project_path:
            if PROJECTS_DIR.exists():
                for project_dir in PROJECTS_DIR.iterdir():
                    if project_dir.is_dir():
                        state_file = project_dir / "state" / f"chunk_{session_id}.json"
                        if state_file.exists():
                            return cls._load_from_file(state_file)

        # Strategy 3: Find most recent state in project
        if project_path and not session_id:
            state_dir = get_state_dir(project_path)
            return cls._load_most_recent_from_dir(state_dir)

        # Strategy 4: Legacy - find most recent state globally
        return cls._load_most_recent_global()

    @classmethod
    def _load_from_file(cls, state_file: Path) -> Optional["ChunkState"]:
        """Load state from a specific file."""
        try:
            data = json.loads(state_file.read_text())
            return cls(**data)
        except (json.JSONDecodeError, TypeError) as e:
            print(f"[CodeContext] Error loading state from {state_file}: {e}")
            return None

    @classmethod
    def _load_most_recent_from_dir(cls, state_dir: Path) -> Optional["ChunkState"]:
        """Find and load the most recent state file in a directory."""
        if not state_dir.exists():
            return None

        state_files = list(state_dir.glob("chunk_*.json"))
        if not state_files:
            return None

        # Sort by modification time, newest first
        state_files.sort(key=lambda f: f.stat().st_mtime, reverse=True)
        return cls._load_from_file(state_files[0])

    @classmethod
    def _load_most_recent_global(cls) -> Optional["ChunkState"]:
        """Find and load the most recent state file across all projects."""
        if not PROJECTS_DIR.exists():
            return None

        all_state_files = []
        for project_dir in PROJECTS_DIR.iterdir():
            if project_dir.is_dir():
                state_dir = project_dir / "state"
                if state_dir.exists():
                    all_state_files.extend(state_dir.glob("chunk_*.json"))

        if not all_state_files:
            return None

        # Sort by modification time, newest first
        all_state_files.sort(key=lambda f: f.stat().st_mtime, reverse=True)
        return cls._load_from_file(all_state_files[0])

    @classmethod
    def clear(cls, session_id: str = None, project_path: str = None) -> None:
        """
        Clear state file(s).

        Args:
            session_id: Session ID to clear (optional)
            project_path: Project path to narrow scope (optional)

        Behavior:
            - If both provided: delete specific state file
            - If only session_id: search and delete across all projects
            - If only project_path: delete all state files for project
            - If neither: no-op (too dangerous to clear all)
        """
        if session_id and project_path:
            # Clear specific state file
            state_dir = get_state_dir(project_path)
            state_file = state_dir / f"chunk_{session_id}.json"
            if state_file.exists():
                state_file.unlink()
                return

        if session_id and not project_path:
            # Search and clear across all projects
            if PROJECTS_DIR.exists():
                for project_dir in PROJECTS_DIR.iterdir():
                    if project_dir.is_dir():
                        state_file = project_dir / "state" / f"chunk_{session_id}.json"
                        if state_file.exists():
                            state_file.unlink()
                            return

        if project_path and not session_id:
            # Clear all state files for project (use with caution)
            state_dir = get_state_dir(project_path)
            for state_file in state_dir.glob("chunk_*.json"):
                state_file.unlink()

    @classmethod
    def exists(cls, session_id: str = None, project_path: str = None) -> bool:
        """Check if state file exists for given session/project."""
        if session_id and project_path:
            state_dir = get_state_dir(project_path)
            state_file = state_dir / f"chunk_{session_id}.json"
            return state_file.exists()

        if session_id:
            # Search all projects
            if PROJECTS_DIR.exists():
                for project_dir in PROJECTS_DIR.iterdir():
                    if project_dir.is_dir():
                        state_file = project_dir / "state" / f"chunk_{session_id}.json"
                        if state_file.exists():
                            return True
            return False

        # Legacy: check if any state exists
        return cls._load_most_recent_global() is not None


def find_jsonl_for_session(session_id: str, project_path: Optional[str] = None) -> Optional[Path]:
    """
    Find JSONL file for a session.

    Args:
        session_id: The session UUID
        project_path: Optional project path to narrow search

    Returns:
        Path to JSONL file, or None if not found
    """
    claude_projects = Path.home() / ".claude" / "projects"

    if not claude_projects.exists():
        return None

    # If we have a project path, encode it and search that folder
    if project_path:
        encoded_path = encode_project_path(project_path)
        project_dir = claude_projects / encoded_path

        if project_dir.exists():
            jsonl_file = project_dir / f"{session_id}.jsonl"
            if jsonl_file.exists():
                return jsonl_file

    # Fallback: search all project folders
    for jsonl_file in claude_projects.rglob(f"{session_id}.jsonl"):
        return jsonl_file

    return None


def get_current_jsonl_path() -> Optional[Path]:
    """
    Try to find the current session's JSONL path.

    This is a best-effort function that looks for recent JSONL files.
    """
    claude_projects = Path.home() / ".claude" / "projects"

    if not claude_projects.exists():
        return None

    # Find most recently modified JSONL
    jsonl_files = list(claude_projects.rglob("*.jsonl"))

    if not jsonl_files:
        return None

    # Sort by modification time, newest first
    jsonl_files.sort(key=lambda f: f.stat().st_mtime, reverse=True)

    return jsonl_files[0]


def cleanup_old_state_files(max_age_hours: int = 24) -> int:
    """
    Clean up old state files that may have been orphaned.

    State files older than max_age_hours are likely from crashed or
    interrupted sessions and can be safely removed.

    Args:
        max_age_hours: Maximum age in hours before cleanup (default: 24)

    Returns:
        Number of files cleaned up
    """
    cleaned = 0
    cutoff = datetime.utcnow().timestamp() - (max_age_hours * 3600)

    if not PROJECTS_DIR.exists():
        return 0

    for project_dir in PROJECTS_DIR.iterdir():
        if project_dir.is_dir():
            state_dir = project_dir / "state"
            if state_dir.exists():
                for state_file in state_dir.glob("chunk_*.json"):
                    if state_file.stat().st_mtime < cutoff:
                        state_file.unlink()
                        cleaned += 1

    return cleaned
