"""Typed data structures shared across the whole pipeline."""
from __future__ import annotations

import hashlib
import re
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from urllib.parse import parse_qs, urlparse, urlunparse

from pydantic import BaseModel, Field

# --------------------------------------------------------------------------
# URL canonicalisation - critical for cross-provider deduplication.
# --------------------------------------------------------------------------
_TRACKING_PREFIXES = ("utm_", "ref_", "mc_", "pk_", "_hs", "vero_", "icid")
_TRACKING_KEYS = {
    "gclid", "fbclid", "msclkid", "igshid", "ref", "referrer", "source",
    "spm", "cmpid", "cmp", "sr_share", "share", "at_medium", "at_campaign",
    "yclid", "dclid", "wt_mc", "trk", "trkCampaign", "_ga",
}


def canonical_url(url: str) -> str:
    """Strip tracking params, normalise host/scheme/trailing slash."""
    if url.startswith("magnet:"):
        # Dedupe magnets on their infohash; the tracker list is noise.
        m = re.search(r"btih:([a-zA-Z0-9]{32,40})", url, re.I)
        return f"magnet:?xt=urn:btih:{m.group(1).lower()}" if m else url.strip()
    try:
        u = urlparse(url.strip())
        if not u.scheme:
            u = u._replace(scheme="https")
        netloc = u.netloc.lower()
        if netloc.startswith("www."):
            netloc = netloc[4:]
        netloc = netloc.replace(":443", "").replace(":80", "")
        keep: List[str] = []
        for k, vals in parse_qs(u.query, keep_blank_values=False).items():
            kl = k.lower()
            if kl in _TRACKING_KEYS or any(kl.startswith(p) for p in _TRACKING_PREFIXES):
                continue
            for v in vals:
                keep.append(f"{k}={v}")
        query = "&".join(sorted(keep))
        path = u.path.rstrip("/") or "/"
        return urlunparse((u.scheme, netloc, path, "", query, ""))
    except Exception:
        return url.strip()


def domain_of(url: str) -> str:
    if url.startswith("magnet:"):
        return "magnet"
    try:
        host = urlparse(url).netloc.lower()
        return host[4:] if host.startswith("www.") else host
    except Exception:
        return ""


def registrable(domain: str) -> str:
    """Cheap eTLD+1 approximation, good enough for authority lookup."""
    parts = domain.split(".")
    if len(parts) <= 2:
        return domain
    two = ".".join(parts[-2:])
    multi = {"co.uk", "co.in", "com.au", "co.jp", "org.uk", "ac.uk", "gov.uk",
             "com.br", "co.za", "com.cn", "ne.jp", "or.jp", "gov.in", "ac.in"}
    if two in multi and len(parts) >= 3:
        return ".".join(parts[-3:])
    return two


# --------------------------------------------------------------------------
# Internal dataclasses
# --------------------------------------------------------------------------
@dataclass
class RawResult:
    """A single row from a single provider's SERP."""
    url: str
    title: str = ""
    snippet: str = ""
    provider: str = ""
    rank: int = 0
    published: Optional[str] = None
    score_hint: float = 0.0
    extra: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.url = canonical_url(self.url)
        self.title = (self.title or "").strip()
        self.snippet = " ".join((self.snippet or "").split())


@dataclass
class Document:
    """A deduplicated candidate, enriched as it flows through the pipeline."""
    url: str
    title: str = ""
    snippet: str = ""
    domain: str = ""
    providers: List[str] = field(default_factory=list)
    ranks: Dict[str, int] = field(default_factory=dict)
    queries: List[str] = field(default_factory=list)
    published: Optional[str] = None
    # Provider payload (magnet, seeders, size, ...) carried through the merge.
    extra: Dict[str, Any] = field(default_factory=dict)

    # populated by the reader
    content: str = ""
    fetched: bool = False
    fetch_error: str = ""
    word_count: int = 0

    # scoring
    rrf: float = 0.0
    bm25: float = 0.0
    authority: float = 0.0
    freshness: float = 0.0
    agreement: float = 0.0
    llm_score: float = 0.0
    score: float = 0.0
    cite_id: int = 0

    @property
    def uid(self) -> str:
        return hashlib.sha1(self.url.encode("utf-8", "ignore")).hexdigest()[:12]

    def best_text(self, limit: int = 1200) -> str:
        return (self.content or self.snippet)[:limit]

    def to_source(self) -> Dict[str, Any]:
        return {
            "id": self.cite_id,
            "url": self.url,
            "title": self.title or self.domain,
            "domain": self.domain,
            "snippet": self.snippet[:400],
            "published": self.published,
            "providers": sorted(set(self.providers)),
            "score": round(self.score, 4),
            "signals": {
                "rrf": round(self.rrf, 4),
                "bm25": round(self.bm25, 4),
                "authority": round(self.authority, 3),
                "freshness": round(self.freshness, 3),
                "agreement": round(self.agreement, 3),
                "llm": round(self.llm_score, 3),
            },
            "read": self.fetched,
            "words": self.word_count,
            **({"download": {
                "magnet": self.extra.get("magnet", ""),
                "size": self.extra.get("size", ""),
                "seeders": self.extra.get("seeders", 0),
                "infohash": self.extra.get("infohash", ""),
                "page": self.extra.get("page", ""),
                "kind": self.extra.get("kind", ""),
            }} if self.extra.get("kind") in ("torrent", "repack") else {}),
        }


@dataclass
class Chunk:
    """A passage of evidence pointing back at its document."""
    text: str
    doc: Document
    position: int = 0
    score: float = 0.0


@dataclass
class Timings:
    t0: float = field(default_factory=time.perf_counter)
    marks: Dict[str, float] = field(default_factory=dict)

    def mark(self, name: str) -> float:
        el = (time.perf_counter() - self.t0) * 1000
        self.marks[name] = round(el, 1)
        return el

    def span(self, name: str, start_ms: float) -> None:
        self.marks[name] = round((time.perf_counter() - self.t0) * 1000 - start_ms, 1)

    def now(self) -> float:
        return (time.perf_counter() - self.t0) * 1000


# --------------------------------------------------------------------------
# Public API models
# --------------------------------------------------------------------------
class Message(BaseModel):
    """One turn of conversation. `sources` echoes what a prior answer cited."""
    role: str = Field(..., description="user | assistant")
    content: str = ""
    sources: List[Dict[str, Any]] = Field(default_factory=list)


class SearchRequest(BaseModel):
    """
    Deliberately tiny public surface.

    `query`, `depth` and `safe` are all a caller ever needs. Everything the
    engine does behind them - which of the 48 sources to fan out to, whether
    to engage the deep/onion tier, which synthesis engine tier to use, how
    many agents to spawn - is derived automatically from the query and these
    three fields.
    """
    query: str = Field(..., min_length=1, max_length=2000,
                       description="Natural-language question")
    depth: str = Field("fast",
                       description="instant | fast | deep | extreme | ultra")
    safe: bool = Field(True, description=(
        "true = filtered sources and synthesis. "
        "false = unrestricted: adult, deep/onion, archive and file indexes "
        "are engaged and synthesis runs on non-refusing engines."))
    stream: bool = Field(False, description=(
        "true = stream the response as Server-Sent Events on this same "
        "endpoint. Omit or set false for a single JSON response."))
    model: str = Field("TeD", description=(
        "Capability tier: TeF (fast+smart) | TeD (deep+smart) | TeM (maximum). "
        "Raw backend model ids are not accepted."))
    messages: List[Message] = Field(default_factory=list, description=(
        "Prior turns. Enables follow-ups like 'give me the full lyrics' - the "
        "engine re-reads the exact source it cited last time instead of "
        "searching the web again."))
    site: Optional[str] = Field(None, description=(
        "Read one site only, e.g. 'genius.com'. Usually inferred from the "
        "query ('lyrics on genius.com') or from conversation history."))
    max_results: int = Field(12, ge=1, le=60, description="Sources returned")
    answer: bool = Field(True, description="Generate a cited synthesized answer")
    read_pages: Optional[bool] = Field(None, description="Override full-page reading")
    include_domains: List[str] = Field(default_factory=list)
    exclude_domains: List[str] = Field(default_factory=list)
    freshness: Optional[str] = Field(None, description="day | week | month | year")
    language: str = Field("en", description="Answer language hint")
    providers: List[str] = Field(default_factory=list, description="Restrict to these providers")
    follow_up: bool = Field(True, description="Allow multi-hop gap-filling searches")
    # Derived internally from `safe` + query intent. Hidden from the public
    # schema so the payload surface stays: query / depth / safe.
    unfiltered: bool = Field(False, exclude=True, json_schema_extra={"readOnly": True})
    cache: bool = True

    model_config = {"json_schema_extra": {
        "example": {"query": "your question", "depth": "deep", "safe": False}}}


class SourceOut(BaseModel):
    download: Optional[Dict[str, Any]] = None
    id: int
    url: str
    title: str
    domain: str
    snippet: str = ""
    published: Optional[str] = None
    providers: List[str] = []
    score: float = 0.0
    signals: Dict[str, float] = {}
    read: bool = False
    words: int = 0


class SearchResponse(BaseModel):
    focus: Optional[Dict[str, Any]] = None
    query: str
    depth: str
    answer: str = ""
    key_points: List[str] = []
    follow_ups: List[str] = []
    sources: List[SourceOut] = []
    sub_queries: List[str] = []
    provider_stats: Dict[str, Any] = {}
    stats: Dict[str, Any] = {}
    tokens: Dict[str, Any] = {}
    timings_ms: Dict[str, float] = {}
    cached: bool = False
    elapsed_ms: float = 0.0