from __future__ import annotations

import json
from datetime import datetime
from pathlib import Path

from pydantic import ValidationError

from .search import SessionSearchIndex
from .models import (
    AssistantRecord,
    Message,
    MessagePayload,
    ProjectInfo,
    Session,
    SessionMetadata,
    SummaryRecord,
    TextBlock,
    UserRecord,
)

CLAUDE_PROJECTS_DIR = Path.home() / ".claude" / "projects"

RELEVANT_TYPES = {"user", "assistant", "summary"}


def discover_projects(base_dir: Path = CLAUDE_PROJECTS_DIR) -> list[ProjectInfo]:
    if not base_dir.exists():
        return []

    projects: list[ProjectInfo] = []
    for entry in sorted(base_dir.iterdir()):
        if not entry.is_dir():
            continue
        sessions = _discover_sessions(entry)
        if not sessions:
            continue

        project_path = sessions[0].project_path or _decode_project_dir(entry.name)
        sessions.sort(key=lambda s: s.timestamp or datetime.min, reverse=True)
        projects.append(ProjectInfo(
            project_dir=entry.name,
            project_path=project_path,
            sessions=sessions,
        ))

    projects.sort(
        key=lambda p: max((s.timestamp for s in p.sessions if s.timestamp), default=datetime.min),
        reverse=True,
    )

    _assign_short_names(projects)
    return projects


def _assign_short_names(projects: list[ProjectInfo]) -> None:
    """Compute shortest unique suffix for each project path."""
    paths = [p.project_path for p in projects]

    def _unique_suffix(path: str, all_paths: list[str]) -> str:
        parts = path.rstrip("/").split("/")
        for depth in range(1, len(parts) + 1):
            suffix = "/".join(parts[-depth:])
            conflicts = sum(1 for p in all_paths if p.rstrip("/").endswith(suffix))
            if conflicts == 1:
                return suffix
        return path

    for project in projects:
        short = _unique_suffix(project.project_path, paths)
        for session in project.sessions:
            session.project_short_name = short


def _discover_sessions(project_dir: Path) -> list[SessionMetadata]:
    sessions: list[SessionMetadata] = []
    for f in project_dir.glob("*.jsonl"):
        meta = _parse_metadata(f, project_dir.name)
        if meta and meta.message_count > 0:
            sessions.append(meta)
    return sessions


def _parse_metadata(file_path: Path, project_dir: str) -> SessionMetadata | None:
    summary = ""
    first_message = ""
    timestamp: datetime | None = None
    last_timestamp: datetime | None = None
    cwd = ""
    git_branch = ""
    message_count = 0
    total_tokens = 0
    seen_request_ids: set[str] = set()
    message_timestamps: list[datetime] = []

    try:
        with open(file_path) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    raw = json.loads(line)
                except json.JSONDecodeError:
                    continue

                rtype = raw.get("type", "")
                if rtype not in RELEVANT_TYPES:
                    continue

                if rtype == "summary" and not summary:
                    summary = raw.get("summary", "")
                elif rtype == "user":
                    message_count += 1
                    if not first_message:
                        content = raw.get("message", {}).get("content", "")
                        if isinstance(content, str) and content:
                            first_message = content[:200]
                    if not timestamp:
                        timestamp = _parse_timestamp(raw.get("timestamp", ""))
                    if not cwd:
                        cwd = raw.get("cwd", "")
                    if not git_branch:
                        git_branch = raw.get("gitBranch", "")
                elif rtype == "assistant":
                    message_count += 1
                    # Extract token usage, deduplicated by requestId
                    request_id = raw.get("requestId", "")
                    if request_id and request_id not in seen_request_ids:
                        seen_request_ids.add(request_id)
                        usage = raw.get("message", {}).get("usage", {})
                        if usage:
                            total_tokens += (
                                usage.get("input_tokens", 0)
                                + usage.get("cache_creation_input_tokens", 0)
                                + usage.get("cache_read_input_tokens", 0)
                                + usage.get("output_tokens", 0)
                            )

                # Track last timestamp and collect for sparkline
                if rtype in ("user", "assistant"):
                    ts = _parse_timestamp(raw.get("timestamp", ""))
                    if ts != datetime.min:
                        last_timestamp = ts
                        message_timestamps.append(ts)
    except (OSError, UnicodeDecodeError):
        return None

    return SessionMetadata(
        session_id=file_path.stem,
        project_dir=project_dir,
        project_path=cwd or _decode_project_dir(project_dir),
        file_path=str(file_path),
        summary=summary,
        first_message_preview=first_message,
        timestamp=timestamp,
        last_timestamp=last_timestamp,
        message_count=message_count,
        git_branch=git_branch,
        total_tokens=total_tokens,
        sparkline_buckets=_compute_sparkline_buckets(message_timestamps),
    )


def load_session(project_dir: str, session_id: str, base_dir: Path = CLAUDE_PROJECTS_DIR) -> Session | None:
    file_path = base_dir / project_dir / f"{session_id}.jsonl"
    if not file_path.exists():
        return None
    return _parse_session(file_path, project_dir)


def _parse_session(file_path: Path, project_dir: str) -> Session:
    summaries: list[str] = []
    ordered: list[UserRecord | AssistantRecord] = []

    with open(file_path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                raw = json.loads(line)
            except json.JSONDecodeError:
                continue

            rtype = raw.get("type", "")
            if rtype == "summary":
                try:
                    rec = SummaryRecord.model_validate(raw)
                    if rec.summary:
                        summaries.append(rec.summary)
                except ValidationError:
                    pass
            elif rtype == "user":
                try:
                    rec = UserRecord.model_validate(raw)
                    ordered.append(rec)
                except ValidationError:
                    pass
            elif rtype == "assistant":
                try:
                    rec = AssistantRecord.model_validate(raw)
                    ordered.append(rec)
                except ValidationError:
                    pass

    messages = _coalesce_messages(ordered)

    # Extract metadata from first user record
    cwd = ""
    first_message = ""
    first_timestamp: datetime | None = None
    git_branch = ""
    for rec in ordered:
        if isinstance(rec, UserRecord):
            if not cwd:
                cwd = rec.cwd
            if not git_branch:
                git_branch = rec.gitBranch
            if not first_message and isinstance(rec.message.content, str) and rec.message.content:
                first_message = rec.message.content[:200]
            if not first_timestamp:
                first_timestamp = _parse_timestamp(rec.timestamp)
            if cwd and first_message and first_timestamp and git_branch:
                break

    # Get last timestamp from the last record
    last_timestamp: datetime | None = None
    for rec in reversed(ordered):
        ts = _parse_timestamp(rec.timestamp)
        if ts != datetime.min:
            last_timestamp = ts
            break

    metadata = SessionMetadata(
        session_id=file_path.stem,
        project_dir=project_dir,
        project_path=cwd or _decode_project_dir(project_dir),
        file_path=str(file_path),
        summary=summaries[0] if summaries else "",
        first_message_preview=first_message,
        timestamp=first_timestamp,
        last_timestamp=last_timestamp,
        message_count=len(messages),
        git_branch=git_branch,
    )

    return Session(metadata=metadata, messages=messages)


def _coalesce_messages(records: list[UserRecord | AssistantRecord]) -> list[Message]:
    messages: list[Message] = []
    i = 0

    while i < len(records):
        rec = records[i]

        if isinstance(rec, UserRecord):
            blocks = _extract_blocks(rec.message)
            if blocks:
                messages.append(Message(
                    role="user",
                    timestamp=_parse_timestamp(rec.timestamp),
                    content_blocks=blocks,
                ))
            i += 1
        elif isinstance(rec, AssistantRecord):
            # Coalesce consecutive assistant records with the same requestId
            request_id = rec.requestId
            blocks: list = []
            timestamp = _parse_timestamp(rec.timestamp)

            while i < len(records):
                r = records[i]
                if not isinstance(r, AssistantRecord) or r.requestId != request_id:
                    break
                blocks.extend(_extract_blocks(r.message))
                i += 1

            if blocks:
                messages.append(Message(
                    role="assistant",
                    timestamp=timestamp,
                    content_blocks=blocks,
                ))
        else:
            i += 1

    return messages


def _extract_blocks(payload: MessagePayload) -> list:
    content = payload.content
    if isinstance(content, str):
        if content.strip():
            return [TextBlock(text=content)]
        return []
    # content is already a list of validated ContentBlock models
    return list(content)


_NUM_SPARKLINE_BUCKETS = 8


def _compute_sparkline_buckets(
    timestamps: list[datetime],
    num_buckets: int = _NUM_SPARKLINE_BUCKETS,
) -> list[int]:
    if len(timestamps) < 2:
        return [len(timestamps)] if timestamps else []

    sorted_ts = sorted(timestamps)
    start = sorted_ts[0]
    end = sorted_ts[-1]
    total_seconds = (end - start).total_seconds()

    if total_seconds <= 0:
        return [len(timestamps)]

    bucket_width = total_seconds / num_buckets
    buckets = [0] * num_buckets

    for ts in sorted_ts:
        offset = (ts - start).total_seconds()
        idx = min(int(offset / bucket_width), num_buckets - 1)
        buckets[idx] += 1

    return buckets


def _decode_project_dir(encoded: str) -> str:
    if encoded.startswith("-"):
        return "/" + encoded[1:].replace("-", "/")
    return encoded


def _parse_timestamp(ts: str) -> datetime:
    if not ts:
        return datetime.min
    try:
        ts = ts.replace("Z", "+00:00")
        return datetime.fromisoformat(ts)
    except (ValueError, TypeError):
        return datetime.min


class SessionIndex:
    def __init__(self):
        self.all_sessions: list[SessionMetadata] = []
        self.project_names: list[str] = []
        self._session_cache: dict[tuple[str, str], Session | None] = {}
        self._search_index: SessionSearchIndex | None = None
        self.loaded: bool = False

    @property
    def search_index(self) -> SessionSearchIndex | None:
        return self._search_index

    def load(self) -> None:
        projects = discover_projects()

        for p in projects:
            self.all_sessions.extend(p.sessions)
        self.all_sessions.sort(key=lambda s: s.timestamp or datetime.min, reverse=True)

        seen: set[str] = set()
        for s in self.all_sessions:
            if s.project_short_name and s.project_short_name not in seen:
                seen.add(s.project_short_name)
                self.project_names.append(s.project_short_name)

        self._search_index = SessionSearchIndex()
        self._search_index.load_or_build(self.all_sessions)

        self.loaded = True

    def get_session(self, project_dir: str, session_id: str) -> Session | None:
        key = (project_dir, session_id)
        if key not in self._session_cache:
            self._session_cache[key] = load_session(project_dir, session_id)
        return self._session_cache[key]


_index = SessionIndex()


def get_index() -> SessionIndex:
    return _index
