"""Thread-safe in-memory + disk-persisted storage for test execution sessions."""

import dataclasses
import json
import threading
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from uuid import uuid4

from maeris_mcp.storage.test_store import _get_sessions_dir


@dataclass
class TestResult:
    """Result of a single test execution attempt."""

    test_file: str
    test_name: str
    status: str  # "passed" | "failed" | "error" | "unfixable"
    attempts: int
    logs: str


@dataclass
class TestExecutionSession:
    """Active session for iterative test execution and auto-fixing."""

    exec_session_id: str
    gen_session_id: str | None
    output_dir: str
    test_files: list[str]  # relative filenames like "my_test.py"
    base_url: str | None = None  # locked base URL — fixes must not change this domain
    current_index: int = 0
    retry_count: int = 0
    max_retries: int = 3
    results: list[TestResult] = field(default_factory=list)
    failed_descriptions: list[str] = field(default_factory=list)
    status: str = "running"  # "running" | "completed"
    created_at: str = ""
    scripts_decision_made: bool = False  # True once save_failing_scripts has been answered


class TestExecutionStore:
    """Thread-safe storage for test execution sessions, persisted to disk."""

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._sessions: dict[str, TestExecutionSession] = {}

    # ------------------------------------------------------------------
    # Disk helpers
    # ------------------------------------------------------------------

    def _session_path(self, exec_session_id: str) -> Path:
        return _get_sessions_dir() / f"exec_{exec_session_id}.json"

    def _save(self, session: TestExecutionSession) -> None:
        """Write session to disk. Never raises — persistence failure is non-fatal."""
        try:
            data = dataclasses.asdict(session)
            self._session_path(session.exec_session_id).write_text(
                json.dumps(data, indent=2), encoding="utf-8"
            )
        except Exception:
            pass

    def _load(self, exec_session_id: str) -> TestExecutionSession | None:
        """Load session from disk. Returns None if not found or corrupt."""
        try:
            path = self._session_path(exec_session_id)
            if not path.exists():
                return None
            data = json.loads(path.read_text(encoding="utf-8"))
            # Reconstruct nested TestResult dataclasses from plain dicts
            data["results"] = [TestResult(**r) for r in data.get("results", [])]
            return TestExecutionSession(**data)
        except Exception:
            return None

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def create_session(
        self,
        output_dir: str,
        test_files: list[str],
        gen_session_id: str | None = None,
        max_retries: int = 3,
        base_url: str | None = None,
    ) -> str:
        """Create a new execution session and return its ID."""
        with self._lock:
            exec_session_id = str(uuid4())
            session = TestExecutionSession(
                exec_session_id=exec_session_id,
                gen_session_id=gen_session_id,
                output_dir=output_dir,
                test_files=test_files,
                base_url=base_url,
                max_retries=max_retries,
                created_at=datetime.now(timezone.utc).isoformat(),
            )
            self._sessions[exec_session_id] = session
            self._save(session)
            return exec_session_id

    def get_session(self, exec_session_id: str) -> TestExecutionSession | None:
        """Retrieve a session by ID.

        Checks memory first; falls back to disk so sessions survive MCP restarts.
        """
        with self._lock:
            session = self._sessions.get(exec_session_id)
            if session:
                return session
            # Try loading from disk (e.g. after MCP restart)
            session = self._load(exec_session_id)
            if session:
                self._sessions[exec_session_id] = session  # warm the in-memory cache
            return session

    def record_result(
        self,
        exec_session_id: str,
        test_file: str,
        test_name: str,
        status: str,
        logs: str,
    ) -> None:
        """Record or update the result for the given test file."""
        with self._lock:
            session = self.get_session(exec_session_id)
            if not session:
                return
            existing = next((r for r in session.results if r.test_file == test_file), None)
            if existing:
                existing.status = status
                existing.attempts += 1
                existing.logs = logs
            else:
                session.results.append(
                    TestResult(
                        test_file=test_file,
                        test_name=test_name,
                        status=status,
                        attempts=session.retry_count + 1,
                        logs=logs,
                    )
                )
            self._save(session)

    def advance(self, exec_session_id: str) -> None:
        """Advance to the next test and reset the retry counter."""
        with self._lock:
            session = self.get_session(exec_session_id)
            if session:
                session.current_index += 1
                session.retry_count = 0
                self._save(session)

    def increment_retry(self, exec_session_id: str) -> None:
        """Increment the retry counter for the current test."""
        with self._lock:
            session = self.get_session(exec_session_id)
            if session:
                session.retry_count += 1
                self._save(session)

    def add_failed_description(self, exec_session_id: str, description: str) -> None:
        """Record a failure description for an unfixable test."""
        with self._lock:
            session = self.get_session(exec_session_id)
            if session:
                session.failed_descriptions.append(description)
                self._save(session)

    def finalize(self, exec_session_id: str) -> None:
        """Mark session as completed."""
        with self._lock:
            session = self.get_session(exec_session_id)
            if session:
                session.status = "completed"
                self._save(session)

    def mark_scripts_resolved(self, exec_session_id: str) -> None:
        """Record that the save_failing_scripts decision has been applied."""
        with self._lock:
            session = self.get_session(exec_session_id)
            if session:
                session.scripts_decision_made = True
                self._save(session)
