"""
Core healing pipeline — orchestrates all healing stages.

Stage execution flow (SMART strategy):
  ┌─────────────────────────┐
  │  0. Try original        │─── found ─────────────────► Return element
  └────────────┬────────────┘
               │ NOT FOUND (quick_timeout_ms)
               ▼
  ┌─────────────────────────┐
  │  1. Check cache         │─── hit ─────► validate ──► Return element
  └────────────┬────────────┘                │ broken
               │ miss                        ▼ (mark miss, continue)
               ▼
  ┌─────────────────────────┐
  │  2. Heuristic healing   │─── found ─────────────────► Cache + Return
  │     (free, ~0ms)        │
  └────────────┬────────────┘
               │ no candidates pass
               ▼
  ┌─────────────────────────┐
  │  3. DOM fuzzy match     │─── found ─────────────────► Cache + Return
  │     (free, ~50ms)       │
  └────────────┬────────────┘
               │ none pass
               ▼
  ┌─────────────────────────┐
  │  4. AI — DOM analysis   │─── found ─────────────────► Cache + Return
  │     (paid, ~1-3s)       │
  └────────────┬────────────┘
               │ all candidates fail
               ▼ (if FULL or VISUAL_FIRST strategy)
  ┌─────────────────────────┐
  │  5. AI — Visual/Shot    │─── found ─────────────────► Cache + Return
  │     (paid, ~2-5s)       │
  └────────────┬────────────┘
               │ all fail
               ▼
         ElementNotFound
"""

from __future__ import annotations

import asyncio
import base64
import logging
import time
from dataclasses import dataclass, field
from typing import Any, List, Optional, Tuple

from playwright.async_api import Locator, Page, TimeoutError as PwTimeout

from playwright_healer.ai_providers import (
    AIHealingResponse,
    HealingCandidate,
    ProviderChain,
)
from playwright_healer.cache import AbstractSelectorCache, CacheEntry, build_cache
from playwright_healer.config import HealerConfig, HealingStrategy
from playwright_healer.heuristic import HeuristicCandidate, HeuristicHealer
from playwright_healer.reporting import (
    ConsoleReporter,
    HealingEvent,
    HTMLReporter,
    JSONReporter,
    SessionReport,
)
from playwright_healer.utils import (
    SelectorType,
    detect_selector_type,
    strip_dom_for_ai,
)

logger = logging.getLogger("playwright_healer.pipeline")


class ElementNotFoundError(Exception):
    """Raised when all healing stages fail to locate the element."""


@dataclass
class HealingResult:
    """Result of a single element lookup attempt."""

    locator: Optional[Locator] = None
    healed_selector: str = ""
    stage: str = "NONE"
    confidence: float = 0.0
    latency_ms: float = 0.0
    tokens_used: int = 0
    provider: str = ""
    model: str = ""
    success: bool = False


class HealingPipeline:
    """
    Orchestrates all healing stages.  One instance per test session;
    shared state is the cache and session report.
    """

    def __init__(self, page: Page, config: HealerConfig, test_name: str = "") -> None:
        self._page = page
        self._cfg = config
        self._test_name = test_name
        self._cache: AbstractSelectorCache = build_cache(config.cache)
        self._chain: Optional[ProviderChain] = (
            ProviderChain(config.providers) if config.providers else None
        )
        self._heuristic = HeuristicHealer()
        self._console = ConsoleReporter()
        self._html_reporter = HTMLReporter()
        self._json_reporter = JSONReporter()
        self._session = SessionReport(
            session_id=f"{int(time.time())}_{id(self)}",
        )

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

    async def find(self, selector: str, description: str = "") -> Locator:
        """
        Main entry point.  Returns a valid Playwright Locator or raises
        ElementNotFoundError.
        """
        start = time.monotonic()
        url = self._page.url or ""
        event = HealingEvent(
            selector=selector,
            description=description,
            url=url,
            test_name=self._test_name,
        )

        try:
            result = await self._run_pipeline(selector, description, url)
        except ElementNotFoundError as exc:
            event.success = False
            event.stage = "NONE"
            event.error = str(exc)[:200]
            event.latency_ms = (time.monotonic() - start) * 1000
            self._session.events.append(event)
            self._console.log_event(event, self._cfg.console_logging)
            raise
        except Exception as exc:
            event.success = False
            event.stage = "NONE"
            event.error = str(exc)[:200]
            event.latency_ms = (time.monotonic() - start) * 1000
            self._session.events.append(event)
            self._console.log_event(event, self._cfg.console_logging)
            raise ElementNotFoundError(
                f"playwright-healer: Could not locate '{selector}' ({description})"
            ) from exc

        if not result.success or result.locator is None:
            event.success = False
            event.stage = result.stage or "NONE"
            event.error = "All healing stages exhausted"
            event.latency_ms = (time.monotonic() - start) * 1000
            self._session.events.append(event)
            self._console.log_event(event, self._cfg.console_logging)
            raise ElementNotFoundError(
                f"playwright-healer: Could not locate '{selector}' ({description})"
            )

        event.healed_selector = result.healed_selector
        event.stage = result.stage
        event.success = result.success
        event.confidence = result.confidence
        event.latency_ms = (time.monotonic() - start) * 1000
        event.tokens_used = result.tokens_used
        event.provider = result.provider
        event.model = result.model

        self._session.events.append(event)
        self._console.log_event(event, self._cfg.console_logging)

        return result.locator

    async def find_all(self, selector: str, description: str = "") -> List[Locator]:
        """Find all matching elements, applying healing if needed."""
        locator = await self.find(selector, description)
        return await locator.all()

    async def is_present(self, selector: str, description: str = "") -> bool:
        """Return True if element exists without raising on not-found."""
        try:
            loc = await self.find(selector, description)
            return await loc.count() > 0
        except (ElementNotFoundError, Exception):
            return False

    async def shutdown(self) -> None:
        """Flush cache, generate reports, print summary."""
        self._session.end_time = time.time()
        self._console.print_summary(self._session, self._cfg.console_logging)
        if self._cfg.report_html:
            path = self._html_reporter.generate(self._session, self._cfg.report_dir)
            logger.info("HTML report: %s", path)
        if self._cfg.report_json:
            path = self._json_reporter.generate(self._session, self._cfg.report_dir)
            logger.info("JSON report: %s", path)

    @property
    def session_report(self) -> SessionReport:
        return self._session

    @property
    def cache(self) -> AbstractSelectorCache:
        return self._cache

    # ------------------------------------------------------------------
    # Internal pipeline
    # ------------------------------------------------------------------

    async def _run_pipeline(
        self,
        selector: str,
        description: str,
        url: str,
    ) -> HealingResult:
        strategy = self._cfg.strategy

        # Stage 0 — Try original selector
        locator = await self._try_locator(selector, timeout_ms=self._cfg.quick_timeout_ms)
        if locator is not None:
            return HealingResult(
                locator=locator,
                healed_selector=selector,
                stage="ORIGINAL",
                confidence=1.0,
                success=True,
            )

        # Stage 1 — Cache lookup
        cache_key = self._cache._cache_key(selector, description)
        cached = self._cache.get(cache_key)
        if cached:
            loc = await self._try_locator_from_entry(cached)
            if loc is not None:
                self._cache.record_hit(cache_key)
                return HealingResult(
                    locator=loc,
                    healed_selector=cached.healed_selector,
                    stage="CACHE",
                    confidence=cached.confidence,
                    success=True,
                )
            else:
                self._cache.record_miss(cache_key)

        if strategy == HealingStrategy.HEURISTIC_ONLY:
            return await self._stage_heuristic(selector, description, url, cache_key)

        if strategy == HealingStrategy.DOM_ONLY:
            return await self._stage_ai_dom(selector, description, url, cache_key)

        if strategy == HealingStrategy.VISUAL_ONLY:
            return await self._stage_ai_visual(selector, description, url, cache_key)

        if strategy == HealingStrategy.PARALLEL:
            return await self._stage_parallel(selector, description, url, cache_key)

        # SMART and FULL: heuristic → DOM fuzzy → AI DOM → (FULL) AI visual
        if self._cfg.enable_heuristic:
            result = await self._stage_heuristic(selector, description, url, cache_key)
            if result.success:
                return result

        if self._cfg.enable_dom_fuzzy:
            result = await self._stage_dom_fuzzy(selector, description, url, cache_key)
            if result.success:
                return result

        result = await self._stage_ai_dom(selector, description, url, cache_key)
        if result.success:
            return result

        if strategy == HealingStrategy.FULL:
            result = await self._stage_ai_visual(selector, description, url, cache_key)
            if result.success:
                return result

        raise ElementNotFoundError(f"All healing stages exhausted for '{selector}'")

    # ------------------------------------------------------------------
    # Stage implementations
    # ------------------------------------------------------------------

    async def _stage_heuristic(
        self,
        selector: str,
        description: str,
        url: str,
        cache_key: str,
    ) -> HealingResult:
        start = time.monotonic()
        try:
            html = await self._page.content()
        except Exception:
            return HealingResult(success=False, stage="HEURISTIC")

        candidates = self._heuristic.heal(selector, description, html)

        for candidate in candidates:
            # Use the candidate's semantic type directly — don't re-detect from string
            loc = await self._try_locator_semantic(
                candidate.selector, candidate.selector_type, timeout_ms=2000
            )
            if loc is not None:
                latency = (time.monotonic() - start) * 1000
                # Store selector in canonical form for the cache
                pw_selector = self._to_playwright_selector(
                    candidate.selector, candidate.selector_type
                )
                entry = self._make_cache_entry(
                    selector, pw_selector, description, url,
                    "HEURISTIC", candidate.confidence,
                )
                self._cache.put(cache_key, entry)
                return HealingResult(
                    locator=loc,
                    healed_selector=pw_selector,
                    stage="HEURISTIC",
                    confidence=candidate.confidence,
                    latency_ms=latency,
                    success=True,
                )
        return HealingResult(success=False, stage="HEURISTIC")

    async def _stage_dom_fuzzy(
        self,
        selector: str,
        description: str,
        url: str,
        cache_key: str,
    ) -> HealingResult:
        """
        Use rapidfuzz to score every candidate element in the DOM against
        the broken selector's attributes, and try the top 5.
        """
        start = time.monotonic()
        from rapidfuzz import fuzz
        from playwright_healer.utils import (
            css_id_from_selector,
            css_classes_from_selector,
            extract_fingerprint,
        )
        try:
            html = await self._page.content()
        except Exception:
            return HealingResult(success=False, stage="DOM_FUZZY")

        try:
            from bs4 import BeautifulSoup
            soup = BeautifulSoup(html, "lxml")
        except Exception:
            return HealingResult(success=False, stage="DOM_FUZZY")

        target_id = css_id_from_selector(selector) or ""
        target_classes = css_classes_from_selector(selector)
        target_desc = description.lower()

        scored = []
        for el in soup.find_all(["input", "button", "a", "select", "textarea",
                                  "label", "div", "span", "h1", "h2", "h3"]):
            fp = extract_fingerprint(el)
            score = 0.0
            if target_id and fp.id:
                score = max(score, fuzz.ratio(target_id, fp.id) / 100.0 * 0.9)
            if target_classes and fp.classes:
                from rapidfuzz import fuzz as f2
                for tc in target_classes:
                    for ec in fp.classes:
                        score = max(score, f2.ratio(tc, ec) / 100.0 * 0.7)
            if target_desc and fp.text:
                score = max(score, fuzz.partial_ratio(target_desc, fp.text.lower()) / 100.0 * 0.6)
            if score >= 0.55:
                scored.append((score, fp, el))

        scored.sort(key=lambda x: -x[0])

        for score, fp, el in scored[:5]:
            from playwright_healer.utils import make_css_from_fingerprint
            css_candidates = make_css_from_fingerprint(fp)
            for css in css_candidates:
                if not css:
                    continue
                loc = await self._try_locator(css, timeout_ms=2000)
                if loc is not None:
                    latency = (time.monotonic() - start) * 1000
                    entry = self._make_cache_entry(
                        selector, css, description, url, "DOM_FUZZY", score
                    )
                    self._cache.put(cache_key, entry)
                    return HealingResult(
                        locator=loc,
                        healed_selector=css,
                        stage="DOM_FUZZY",
                        confidence=score,
                        latency_ms=latency,
                        success=True,
                    )

        return HealingResult(success=False, stage="DOM_FUZZY")

    async def _stage_ai_dom(
        self,
        selector: str,
        description: str,
        url: str,
        cache_key: str,
    ) -> HealingResult:
        if not self._chain:
            return HealingResult(success=False, stage="AI_DOM")

        start = time.monotonic()
        try:
            html = await self._page.content()
            dom = strip_dom_for_ai(html)
            response = await self._chain.heal_with_dom(selector, description, dom, url)
        except Exception as exc:
            logger.warning("AI DOM healing failed: %s", exc)
            return HealingResult(success=False, stage="AI_DOM")

        for candidate in response.candidates:
            pw_selector = self._to_playwright_selector(
                candidate.selector, candidate.selector_type
            )
            loc = await self._try_locator(pw_selector, timeout_ms=3000)
            if loc is not None:
                latency = (time.monotonic() - start) * 1000
                entry = self._make_cache_entry(
                    selector, pw_selector, description, url,
                    "AI_DOM", candidate.confidence,
                )
                if self._cfg.adaptive_ttl:
                    entry.expires_at = time.time() + self._cache._adaptive_ttl(entry)
                else:
                    entry.expires_at = time.time() + self._cache._ttl_seconds(self._cfg.cache)
                self._cache.put(cache_key, entry)
                return HealingResult(
                    locator=loc,
                    healed_selector=pw_selector,
                    stage="AI_DOM",
                    confidence=candidate.confidence,
                    latency_ms=latency,
                    tokens_used=response.tokens_used,
                    provider=response.provider,
                    model=response.model,
                    success=True,
                )
        return HealingResult(success=False, stage="AI_DOM")

    async def _stage_ai_visual(
        self,
        selector: str,
        description: str,
        url: str,
        cache_key: str,
    ) -> HealingResult:
        if not self._chain or not self._chain.has_vision():
            return HealingResult(success=False, stage="AI_VISUAL")

        start = time.monotonic()
        try:
            screenshot_bytes = await self._page.screenshot(
                type="png", full_page=False
            )
            screenshot_b64 = base64.b64encode(screenshot_bytes).decode()
            response = await self._chain.heal_with_vision(
                selector, description, screenshot_b64, url
            )
        except Exception as exc:
            logger.warning("AI visual healing failed: %s", exc)
            return HealingResult(success=False, stage="AI_VISUAL")

        for candidate in response.candidates:
            pw_selector = self._to_playwright_selector(
                candidate.selector, candidate.selector_type
            )
            loc = await self._try_locator(pw_selector, timeout_ms=3000)
            if loc is not None:
                latency = (time.monotonic() - start) * 1000
                entry = self._make_cache_entry(
                    selector, pw_selector, description, url,
                    "AI_VISUAL", candidate.confidence,
                )
                self._cache.put(cache_key, entry)
                return HealingResult(
                    locator=loc,
                    healed_selector=pw_selector,
                    stage="AI_VISUAL",
                    confidence=candidate.confidence,
                    latency_ms=latency,
                    tokens_used=response.tokens_used,
                    provider=response.provider,
                    model=response.model,
                    success=True,
                )
        return HealingResult(success=False, stage="AI_VISUAL")

    async def _stage_parallel(
        self,
        selector: str,
        description: str,
        url: str,
        cache_key: str,
    ) -> HealingResult:
        """Run DOM and Visual AI stages concurrently; use highest-confidence hit."""
        results = await asyncio.gather(
            self._stage_ai_dom(selector, description, url, cache_key),
            self._stage_ai_visual(selector, description, url, cache_key),
            return_exceptions=True,
        )
        valid = [r for r in results if isinstance(r, HealingResult) and r.success]
        if not valid:
            raise ElementNotFoundError(f"Parallel AI healing exhausted for '{selector}'")
        best = max(valid, key=lambda r: r.confidence)
        return best

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    async def _try_locator(
        self, selector: str, timeout_ms: int = 500
    ) -> Optional[Locator]:
        """Try a selector (auto-detect type) and return a Locator if found, else None."""
        try:
            locator = self._selector_to_locator(selector)
            await locator.wait_for(state="attached", timeout=timeout_ms)
            count = await locator.count()
            return locator if count > 0 else None
        except (PwTimeout, Exception):
            return None

    async def _try_locator_semantic(
        self, selector: str, selector_type: str, timeout_ms: int = 2000
    ) -> Optional[Locator]:
        """Try a selector using the EXPLICIT semantic type (bypasses auto-detect)."""
        try:
            locator = self._semantic_type_to_locator(selector, selector_type)
            await locator.wait_for(state="attached", timeout=timeout_ms)
            count = await locator.count()
            return locator if count > 0 else None
        except (PwTimeout, Exception):
            return None

    def _semantic_type_to_locator(self, selector: str, selector_type: str) -> Locator:
        """Build a Playwright Locator from an explicit selector type."""
        t = selector_type.lower()
        if t == "text":
            return self._page.get_by_text(selector, exact=False)
        if t == "placeholder":
            return self._page.get_by_placeholder(selector)
        if t == "label":
            return self._page.get_by_label(selector)
        if t == "testid":
            return self._page.get_by_test_id(selector)
        if t == "alt":
            return self._page.get_by_alt_text(selector)  # type: ignore[attr-defined]
        if t == "title":
            return self._page.get_by_title(selector)
        if t == "role":
            parts = selector.split("::", 1)
            role = parts[0].strip()
            name = parts[1].strip() if len(parts) > 1 else ""
            if name:
                return self._page.get_by_role(role, name=name)  # type: ignore[arg-type]
            return self._page.get_by_role(role)  # type: ignore[arg-type]
        # CSS / XPath default
        return self._page.locator(selector)

    async def _try_locator_from_entry(self, entry: CacheEntry) -> Optional[Locator]:
        return await self._try_locator(entry.healed_selector, timeout_ms=2000)

    def _selector_to_locator(self, selector: str) -> Locator:
        """Convert a canonical selector string to a Playwright Locator."""
        sel_type = detect_selector_type(selector)

        if sel_type == SelectorType.ROLE:
            # Format: ROLE::NAME
            parts = selector.split("::", 1)
            role = parts[0].strip()
            name = parts[1].strip() if len(parts) > 1 else ""
            if name:
                return self._page.get_by_role(role, name=name)  # type: ignore[arg-type]
            return self._page.get_by_role(role)  # type: ignore[arg-type]

        if sel_type == SelectorType.TEXT:
            return self._page.get_by_text(selector, exact=False)

        if sel_type == SelectorType.PLACEHOLDER:
            return self._page.get_by_placeholder(selector)

        if sel_type == SelectorType.LABEL:
            return self._page.get_by_label(selector)

        if sel_type == SelectorType.TEST_ID:
            return self._page.get_by_test_id(selector)

        if sel_type == SelectorType.XPATH:
            return self._page.locator(selector)

        # CSS (default)
        return self._page.locator(selector)

    def _to_playwright_selector(self, selector: str, selector_type: str) -> str:
        """Normalise a candidate selector into canonical form used by this pipeline."""
        t = selector_type.lower()

        # Role selectors: ensure ROLE::NAME format
        if t == "role":
            if "::" not in selector:
                return selector  # already normalised or just a role name
            return selector

        # Everything else passes through as-is
        return selector

    def _make_cache_entry(
        self,
        original: str,
        healed: str,
        description: str,
        url: str,
        stage: str,
        confidence: float,
    ) -> CacheEntry:
        import time as _time
        return CacheEntry(
            original_selector=original,
            healed_selector=healed,
            description=description,
            url_pattern=url,
            stage=stage,
            confidence=confidence,
            expires_at=_time.time() + self._cfg.cache.ttl_hours * 3600,
        )
