from pydantic import BaseModel, HttpUrl, Field
from typing import List, Optional, Tuple, Literal, Dict
from datetime import datetime

class SourceMeta(BaseModel):
    jurisdiction: Literal["HK", "EU", "ICAO", "SESAR", "ASTM", "ISO", "OTHER"]
    authority: str
    collection: str
    canonical_id: str
    title: str
    lang: str = "en"
    version: Optional[str] = None
    effective_date: Optional[datetime] = None
    publication_date: Optional[datetime] = None
    retrieved_at: datetime = Field(default_factory=datetime.utcnow)
    url: HttpUrl
    media_type: Literal["application/pdf", "text/html"]
    etag: Optional[str] = None
    last_modified: Optional[str] = None
    sha256: Optional[str] = None
    pages: Optional[int] = None
    notes: Optional[str] = None

class Node(BaseModel):
    node_id: str
    path: List[str]
    label: str
    title: Optional[str] = None
    page_span: Optional[Tuple[int,int]] = None
    text: str
    html: Optional[str] = None
    numbering: Optional[str] = None
    children: List["Node"] = Field(default_factory=list)
    node_hash: Optional[str] = None
Node.model_rebuild()

class Chunk(BaseModel):
    source: SourceMeta
    node_id: str
    chunk_id: str
    path: List[str]
    text: str
    html: Optional[str] = None
    page_span: Optional[Tuple[int,int]] = None
    token_count: int
    heading_breadcrumb: Optional[str] = None
    numbering: Optional[str] = None
    citations_in: Optional[List[Dict[str, str]]] = None
    tags: Optional[List[str]] = None
    concepts: Optional[List[str]] = None  # concept_ids linked

def iter_structure_chunks(node, max_tokens=1024, min_tokens=200):
    # yield chunks aligned to node/subnode; merge/split to stay within bounds
    if not node.children:
        toks = count_tokens(node.text)
        if toks <= max_tokens:
            yield node, [node.text]
        else:
            for part in split_by_subparagraph(node.text, max_tokens):
                yield node, [part]
        return
    # internal node: recurse and optionally merge tiny leaves
    buffer, buf_tokens = [], 0
    for child in node.children:
        for leaf, parts in iter_structure_chunks(child, max_tokens, min_tokens):
            for p in parts:
                t = count_tokens(p)
                if buf_tokens + t < max_tokens:
                    buffer.append(p); buf_tokens += t
                else:
                    yield child, buffer; buffer, buf_tokens = [p], t
    if buffer:
        yield node, buffer

import hashlib, json
def stable_id(*parts):
    s = "||".join(p for p in parts if p)
    return hashlib.sha256(s.encode("utf-8")).hexdigest()[:16]

node_id = stable_id(source.sha256 or source.url, "/".join(path), numbering or "")
chunk_id = stable_id(node_id, str(idx))

PATTERNS = {
  "REMOTE_ID_BCAST": [r"\b(remote|electronic)\s+id(entification)?\b", r"\bbroadcast\b"],
  "GEOAWARENESS": [r"\bgeo[-\s]?awareness\b|\bno[-\s]?fly\b|\brestricted\s+area\b"],
  "U_SPACE_NETWORK_RID": [r"\bnetwork\s+remote\s+id(entification)?\b"]
}
def auto_concepts(text):
    hits = []
    low = text.lower()
    for cid, regexes in PATTERNS.items():
        if sum(bool(re.search(rx, low)) for rx in regexes) >= 2:
            hits.append(cid)
    return hits