from __future__ import annotations

import json
from pathlib import Path

from .models import ChatEntry, StoreIndex, TaskState, TaskSummary


class StateStore:
    def __init__(self, path: Path) -> None:
        self.path = path
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.tasks_dir = self.path.parent / "tasks"
        self.tasks_dir.mkdir(parents=True, exist_ok=True)

    def load(self, task_id: str | None = None) -> TaskState:
        self._ensure_initialized()
        target_task_id = task_id or self.load_index().active_task_id
        task_path = self.tasks_dir / f"{target_task_id}.json"
        data = json.loads(task_path.read_text(encoding="utf-8"))
        state = TaskState.model_validate(data)
        if self._backfill_chat_log(state):
            self.save(state, task_id=target_task_id)
        return state

    def save(self, state: TaskState, task_id: str | None = None) -> None:
        self._ensure_initialized()
        target_task_id = task_id or state.task_id
        state.task_id = target_task_id
        task_path = self.tasks_dir / f"{target_task_id}.json"
        task_path.write_text(
            json.dumps(state.model_dump(mode="json"), ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        self._update_summary(state)

    def load_index(self) -> StoreIndex:
        self._ensure_initialized()
        data = json.loads(self.path.read_text(encoding="utf-8"))
        return StoreIndex.model_validate(data)

    def save_index(self, index: StoreIndex) -> None:
        self.path.write_text(
            json.dumps(index.model_dump(mode="json"), ensure_ascii=False, indent=2),
            encoding="utf-8",
        )

    def list_tasks(self) -> list[TaskSummary]:
        return self.load_index().tasks

    def set_active_task(self, task_id: str) -> None:
        index = self.load_index()
        index.active_task_id = task_id
        self.save_index(index)

    def create_task(self, task_id: str, title: str, workdir: str | None = None) -> TaskState:
        state = TaskState(task_id=task_id, title=title, workdir=workdir)
        self.save(state, task_id=task_id)
        self.set_active_task(task_id)
        return state

    def _ensure_initialized(self) -> None:
        if not self.path.exists():
            state = TaskState()
            self.save_index(StoreIndex(active_task_id=state.task_id, tasks=[]))
            self.save(state, task_id=state.task_id)
            return

        try:
            data = json.loads(self.path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            data = {}

        if "active_task_id" in data and "tasks" in data:
            return

        legacy_state = TaskState.model_validate(data or {})
        self.save_index(StoreIndex(active_task_id=legacy_state.task_id, tasks=[]))
        self.save(legacy_state, task_id=legacy_state.task_id)

    def _update_summary(self, state: TaskState) -> None:
        index = self.load_index()
        summary = TaskSummary(
            task_id=state.task_id,
            title=state.title,
            status=state.status,
            updated_at=state.updated_at,
            pending_items=sum(1 for note in state.notes if note.processed_at is None),
            completed_steps=state.completed_steps,
        )
        tasks = [task for task in index.tasks if task.task_id != state.task_id]
        tasks.append(summary)
        tasks.sort(key=lambda item: item.updated_at, reverse=True)
        index.tasks = tasks
        if not index.active_task_id:
            index.active_task_id = state.task_id
        self.save_index(index)

    def _backfill_chat_log(self, state: TaskState) -> bool:
        if state.chat_log:
            return False

        entries: list[ChatEntry] = []
        for note in state.notes:
            if note.author in {"mobile", "telegram"}:
                entries.append(ChatEntry(created_at=note.created_at, author=note.author, text=note.text))
                if note.result_summary:
                    entries.append(
                        ChatEntry(
                            created_at=note.processed_at or note.created_at,
                            author="assistant",
                            text=note.result_summary,
                        )
                    )
                continue

            entries.append(
                ChatEntry(
                    created_at=note.processed_at or note.created_at,
                    author="assistant",
                    text=note.result_summary or note.text,
                )
            )

        for record in state.command_history:
            entries.append(
                ChatEntry(
                    created_at=record.created_at,
                    author="assistant",
                    text=f"已執行命令：{record.command}\n退出碼：{record.exit_code}\n{record.output_preview}",
                )
            )

        if not entries:
            return False

        entries.sort(key=lambda item: item.created_at)
        state.chat_log = entries[-120:]
        return True
