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

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

from maeris_mcp.types.test_generation import TestCase


_STEP_FIELD_ALIASES: dict[str, str] = {
    "action": "action_taken",
    "selector": "css_selector",
}


def _normalize_step(step: dict) -> dict:
    """Normalize step field names using known aliases.

    For each alias key: if the alias key exists AND the canonical key does NOT exist,
    rename it (delete old key, add canonical key). If both exist, keep canonical and
    remove alias.
    """
    if not isinstance(step, dict):
        return step
    result = dict(step)
    for alias, canonical in _STEP_FIELD_ALIASES.items():
        if alias in result:
            if canonical not in result:
                result[canonical] = result.pop(alias)
            else:
                del result[alias]
    return result


@dataclass
class TestGenerationSession:
    """Active session for LLM-guided test case generation."""

    session_id: str
    flow_description: str
    target_path: str
    scan_root: str
    base_for_rel: str
    file_list: list[str]
    base_url: str | None = None
    test_cases: list[dict] = field(default_factory=list)
    flow_analysis: dict | None = None
    created_at: str = ""
    status: str = "analyzing"  # "analyzing" | "generating" | "completed"
    output_dir: str | None = None
    # Stores structured flow metadata submitted via setup_ui_automation
    identified_flows: list[dict] = field(default_factory=list)
    # Credentials provided by the user for login/auth test cases
    test_credentials: dict | None = None
    # Auth configuration identified by the LLM during Phase 1 scanning.
    # Includes login selectors and the list of routes/actions that require auth.
    auth_config: dict | None = None
    # Accumulated scan notes: structured observations the LLM submits during Phase 1.
    # Persists across batches so the LLM always has a compact summary of what it learned.
    scan_notes: dict = field(default_factory=dict)
    # Approved plan from Phase 1.5 (scan -> plan -> generate).
    test_plan: dict | None = None
    # Multi-repo context: maps absolute file path -> display key shown to Claude.
    # e.g. "/home/user/backend-api/routes/auth.py" -> "[backend-api]/routes/auth.py"
    # Linked files are prepended to the file list during Phase 1 scanning so Claude
    # sees API routes/schemas/models from sibling repos as context.
    linked_file_map: dict = field(default_factory=dict)
    # Per-call token usage log. Each entry records estimated token counts for that
    # generate_tests invocation so the user can see where tokens are spent.
    token_log: list[dict] = field(default_factory=list)
    # Relative paths of files that changed since the last scan (from scan_cache.get_changed_files).
    # Stored on the session so the completeness gate can mark cached page_elements
    # as stale for routes whose source file changed.
    changed_files: list[str] = field(default_factory=list)
    # Real selector values extracted server-side from source files during Phase 1.
    # Accumulated across all batches. Used to cross-reference page_elements submitted
    # by the LLM and flag any selector that was never seen in source (hallucinated).
    source_selectors: list[str] = field(default_factory=list)
    # Number of times the completeness gate has blocked this session's plan submission.
    # After MAX_GATE_REJECTIONS the gate approves with warnings to prevent infinite loops.
    gate_rejections: int = 0
    # Number of times the flow-coverage gate has blocked plan submission.
    # After MAX_COVERAGE_REJECTIONS the plan proceeds with warnings.
    coverage_rejections: int = 0
    # Number of times the step-reachability gate has rejected submitted test cases.
    # After MAX_REACHABILITY_REJECTIONS the gate approves with warnings to prevent loops.
    reachability_rejections: int = 0
    # Number of times _validate_test_plan hard errors have blocked plan submission.
    # After MAX_VALIDATION_REJECTIONS the validation errors are demoted to warnings
    # to prevent infinite loops when the LLM cannot satisfy the keyword checks.
    validation_rejections: int = 0
    # Number of times Gate 2 (Steps vs Plan) has rejected submitted test_cases.
    # After MAX_PLAN_REJECTIONS the gate approves with warnings to prevent loops.
    plan_mismatch_rejections: int = 0
    # True once Phase 1 scanning has completed (hasMore=false at least once).
    scan_complete: bool = False


def _get_sessions_dir() -> Path:
    """Return the directory where session files are persisted."""
    config_dir = Path(os.getenv("MAERIS_CONFIG_DIR") or (Path.home() / ".maeris"))
    sessions_dir = config_dir / "sessions"
    sessions_dir.mkdir(parents=True, exist_ok=True)
    return sessions_dir


class TestGenerationStore:
    """Thread-safe storage for test generation sessions, persisted to disk."""

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

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

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

    def _save(self, session: TestGenerationSession) -> None:
        """Write session to disk. Never raises — persistence failure is non-fatal."""
        try:
            # Sanitise scan_notes before writing — if it somehow became a
            # non-dict at runtime, fix it now so the on-disk copy stays clean.
            if not isinstance(session.scan_notes, dict):
                session.scan_notes = {}
            data = dataclasses.asdict(session)
            self._session_path(session.session_id).write_text(
                json.dumps(data, indent=2), encoding="utf-8"
            )
        except Exception:
            pass

    def _load(self, session_id: str) -> TestGenerationSession | None:
        """Load session from disk. Returns None if not found or corrupt."""
        try:
            path = self._session_path(session_id)
            if not path.exists():
                return None
            data = json.loads(path.read_text(encoding="utf-8"))
            # Ensure scan_notes is always a dict after deserialization.
            # It can become a string if a previous save/load cycle serialised
            # the dict as a JSON string instead of a plain object.
            sn = data.get("scan_notes")
            if isinstance(sn, str):
                try:
                    data["scan_notes"] = json.loads(sn)
                except (json.JSONDecodeError, TypeError):
                    data["scan_notes"] = {}
            elif not isinstance(sn, dict):
                data["scan_notes"] = {}
            # Provide defaults for fields added after initial release so old sessions load cleanly.
            data.setdefault("changed_files", [])
            data.setdefault("source_selectors", [])
            data.setdefault("coverage_rejections", 0)
            data.setdefault("validation_rejections", 0)
            data.setdefault("scan_complete", False)
            return TestGenerationSession(**data)
        except Exception:
            return None

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

    def create_session(
        self,
        flow_description: str,
        target_path: str,
        scan_root: str,
        base_for_rel: str,
        file_list: list[str],
        base_url: str | None = None,
        linked_file_map: dict | None = None,
    ) -> str:
        """Create a new test generation session and return its session_id."""
        with self._lock:
            session_id = str(uuid4())
            session = TestGenerationSession(
                session_id=session_id,
                flow_description=flow_description,
                target_path=target_path,
                scan_root=scan_root,
                base_for_rel=base_for_rel,
                file_list=file_list,
                base_url=base_url,
                linked_file_map=linked_file_map or {},
                created_at=datetime.now(timezone.utc).isoformat(),
            )
            self._sessions[session_id] = session
            self._save(session)
            return session_id

    def get_session(self, session_id: str) -> TestGenerationSession | None:
        """Retrieve a test generation session by ID.

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

    def add_test_cases(self, session_id: str, test_cases: list[dict]) -> dict[str, int]:
        """Add test cases to a session. Normalizes step field name aliases before storing."""
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return {"added": 0}

            normalized: list[dict] = []
            for tc in test_cases:
                if isinstance(tc, dict):
                    steps = tc.get("steps")
                    if isinstance(steps, list):
                        tc = {**tc, "steps": [_normalize_step(s) for s in steps]}
                normalized.append(tc)

            added = len(normalized)
            session.test_cases.extend(normalized)
            self._save(session)
            return {"added": added, "total": len(session.test_cases)}

    def remove_last_test_cases(self, session_id: str, count: int) -> None:
        """Remove the last `count` test cases added to a session (used to roll back a rejected batch)."""
        with self._lock:
            session = self.get_session(session_id)
            if not session or count <= 0:
                return
            if count >= len(session.test_cases):
                session.test_cases = []
            else:
                session.test_cases = session.test_cases[:-count]
            self._save(session)

    def set_credentials(self, session_id: str, credentials: dict) -> bool:
        """Store user-provided test credentials on the session."""
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return False
            session.test_credentials = credentials
            self._save(session)
            return True

    def add_identified_flows(self, session_id: str, flows: list[dict]) -> dict[str, int]:
        """Accumulate structured flow metadata and flatten test_cases. Enforces max 7 flows."""
        MAX_FLOWS = 7
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return {"added": 0}
            added = 0
            new_flat_tcs: list[dict] = []
            for flow in flows:
                name = flow.get("name")
                existing = next(
                    (f for f in session.identified_flows if f.get("name") == name), None
                )
                if existing:
                    # Merge new test_cases into existing flow entry
                    existing.setdefault("test_cases", []).extend(flow.get("test_cases", []))
                elif len(session.identified_flows) < MAX_FLOWS:
                    session.identified_flows.append(flow)
                    added += 1
                new_flat_tcs.extend(flow.get("test_cases", []))
            # Populate session.test_cases for push_tests_to_maeris compatibility
            if new_flat_tcs:
                self.add_test_cases(session_id, new_flat_tcs)
            else:
                self._save(session)
            return {"added": added, "total": len(session.identified_flows)}

    def set_auth_config(self, session_id: str, auth_config: dict) -> bool:
        """Store auth configuration identified by the LLM during scanning."""
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return False
            session.auth_config = auth_config
            self._save(session)
            return True

    def update_scan_notes(self, session_id: str, new_notes: dict) -> dict:
        """Merge new scan_notes into the session's accumulated notes.

        List fields (routes, navigation, forms) are appended and deduplicated.
        Dict fields (auth, key_selectors) are shallow-merged.
        page_elements are merged per-page with element-level dedup by selector.
        page_graph edges are deduplicated by (from_page, to_page) tuple.
        Returns the full accumulated scan_notes.
        """
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return {}

            existing = session.scan_notes if isinstance(session.scan_notes, dict) else {}

            # Merge list fields by appending and deduplicating
            for list_key, dedup_field in [
                ("routes", "path"),
                ("navigation", "component"),
                ("forms", "form_name"),
                ("conditional_ui", "element_description"),
                ("server_redirects", "source"),
                ("auth_guards", "component"),
            ]:
                old_items = existing.get(list_key, [])
                new_items = new_notes.get(list_key, [])
                if new_items:
                    seen = set()
                    for item in old_items:
                        if not isinstance(item, dict):
                            continue
                        dedup_val = item.get(dedup_field)
                        if list_key == "forms" and not dedup_val:
                            dedup_val = item.get("purpose")
                        if dedup_val:
                            seen.add(dedup_val)
                    for item in new_items:
                        if not isinstance(item, dict):
                            continue
                        key = item.get(dedup_field)
                        if list_key == "forms" and not key:
                            key = item.get("purpose")
                        if key and key not in seen:
                            old_items.append(item)
                            seen.add(key)
                        elif not key:
                            old_items.append(item)
                    existing[list_key] = old_items

            # Merge dict fields by shallow merge
            for dict_key in ("auth", "key_selectors", "app_domain_info", "error_messages", "success_indicators", "navigation_tree"):
                old_dict = existing.get(dict_key, {})
                new_dict = new_notes.get(dict_key, {})
                if new_dict:
                    if isinstance(old_dict, dict) and isinstance(new_dict, dict):
                        old_dict.update(new_dict)
                    else:
                        old_dict = new_dict
                    existing[dict_key] = old_dict

            # Merge page_elements: group by page, dedup elements by selector
            new_page_elements = new_notes.get("page_elements", [])
            if new_page_elements:
                old_page_elements = existing.get("page_elements", [])
                page_index: dict[str, dict] = {}
                for entry in old_page_elements:
                    if not isinstance(entry, dict):
                        continue
                    page = entry.get("page", "")
                    if page:
                        page_index[page] = entry
                for entry in new_page_elements:
                    if not isinstance(entry, dict):
                        continue
                    page = entry.get("page", "")
                    if not page:
                        continue
                    if page in page_index:
                        # Merge elements into existing page entry
                        old_els = page_index[page].get("elements", [])
                        seen_sels = {
                            el.get("selector") for el in old_els
                            if isinstance(el, dict) and el.get("selector")
                        }
                        for el in entry.get("elements", []):
                            if not isinstance(el, dict):
                                continue
                            sel = el.get("selector")
                            if sel and sel not in seen_sels:
                                old_els.append(el)
                                seen_sels.add(sel)
                        page_index[page]["elements"] = old_els
                    else:
                        page_index[page] = entry
                existing["page_elements"] = list(page_index.values())

            # Merge page_graph: dedup edges by (from_page, to_page)
            new_edges = new_notes.get("page_graph", [])
            if new_edges:
                old_edges = existing.get("page_graph", [])
                seen_edges = {
                    (e.get("from_page", ""), e.get("to_page", ""))
                    for e in old_edges if isinstance(e, dict)
                }
                for edge in new_edges:
                    if not isinstance(edge, dict):
                        continue
                    key = (edge.get("from_page", ""), edge.get("to_page", ""))
                    if key not in seen_edges:
                        old_edges.append(edge)
                        seen_edges.add(key)
                existing["page_graph"] = old_edges

            # Merge nav_graph: dedup edges by (from_page, to_page, trigger)
            new_nav_edges = new_notes.get("nav_graph", [])
            if new_nav_edges:
                old_nav_edges = existing.get("nav_graph", [])
                seen_nav_edges = {
                    (e.get("from_page", ""), e.get("to_page", ""), e.get("trigger", ""))
                    for e in old_nav_edges if isinstance(e, dict)
                }
                for edge in new_nav_edges:
                    if not isinstance(edge, dict):
                        continue
                    key = (
                        edge.get("from_page", ""),
                        edge.get("to_page", ""),
                        edge.get("trigger", ""),
                    )
                    if key not in seen_nav_edges:
                        old_nav_edges.append(edge)
                        seen_nav_edges.add(key)
                existing["nav_graph"] = old_nav_edges

            # Merge route access lists while deduplicating.
            for list_key in ("blocked_routes", "accessible_routes", "no_elements_routes"):
                old_routes = existing.get(list_key, [])
                new_routes = new_notes.get(list_key, [])
                if isinstance(new_routes, list):
                    for route in new_routes:
                        if isinstance(route, str) and route not in old_routes:
                            old_routes.append(route)
                    existing[list_key] = old_routes

            session.scan_notes = existing
            self._save(session)
            return existing

    def update_changed_files(self, session_id: str, changed_files: list[str]) -> None:
        """Store the list of changed files (relative paths) on the session.

        Used by the completeness gate to identify stale cached page_elements.
        Never raises.
        """
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return
            session.changed_files = list(changed_files)
            self._save(session)

    def update_source_selectors(self, session_id: str, selectors: set[str]) -> None:
        """Merge newly extracted source selector values into the session's index.

        Called after each file-read batch so the index grows across calls.
        Never raises.
        """
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return
            existing = set(session.source_selectors)
            existing.update(selectors)
            session.source_selectors = sorted(existing)
            self._save(session)

    def increment_gate_rejections(self, session_id: str) -> int:
        """Increment the gate rejection counter and return the new count. Never raises."""
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return 0
            session.gate_rejections = (session.gate_rejections or 0) + 1
            self._save(session)
            return session.gate_rejections

    def increment_coverage_rejections(self, session_id: str) -> int:
        """Increment the flow-coverage rejection counter and return the new count. Never raises."""
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return 0
            session.coverage_rejections = (session.coverage_rejections or 0) + 1
            self._save(session)
            return session.coverage_rejections

    def increment_validation_rejections(self, session_id: str) -> int:
        """Increment the plan-validation hard-error rejection counter. Never raises."""
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return 0
            session.validation_rejections = (session.validation_rejections or 0) + 1
            self._save(session)
            return session.validation_rejections

    def set_scan_complete(self, session_id: str, complete: bool = True) -> bool:
        """Mark scan completion for a session."""
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return False
            session.scan_complete = bool(complete)
            self._save(session)
            return True

    def append_token_log(self, session_id: str, entry: dict) -> None:
        """Append a token-usage entry to the session log. Never raises."""
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return
            session.token_log.append(entry)
            self._save(session)

    def set_test_plan(self, session_id: str, test_plan: dict) -> bool:
        """Store the approved Phase 1.5 test plan on the session."""
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return False
            session.test_plan = test_plan
            if session.status == "analyzing":
                session.status = "generating"
            self._save(session)
            return True

    def update_base_url(self, session_id: str, base_url: str) -> bool:
        """Update the base URL for a session."""
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return False
            session.base_url = base_url
            self._save(session)
            return True

    def finalize_session(
        self,
        session_id: str,
        output_dir: str,
    ) -> TestGenerationSession | None:
        """Finalize a test generation session after scripts are generated."""
        with self._lock:
            session = self.get_session(session_id)
            if not session:
                return None

            session.status = "completed"
            session.output_dir = output_dir
            self._save(session)
            return session

    def get_recent(self, limit: int = 10) -> list[TestGenerationSession]:
        """Get the most recent test generation sessions."""
        with self._lock:
            sessions = sorted(
                self._sessions.values(),
                key=lambda s: s.created_at,
                reverse=True,
            )
            return sessions[:limit]

    def delete(self, session_id: str) -> bool:
        """Delete a test generation session by ID."""
        with self._lock:
            removed = self._sessions.pop(session_id, None) is not None
            path = self._session_path(session_id)
            if path.exists():
                try:
                    path.unlink()
                except Exception:
                    pass
            return removed

    def clear(self) -> int:
        """Clear all in-memory test generation sessions, return count deleted."""
        with self._lock:
            count = len(self._sessions)
            self._sessions.clear()
            return count
