from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
import re
import struct

import numpy as np


DAY_RECORD_SIZE = 32
DAY_STRUCT = struct.Struct("<IIIIIfII")
TNF_HEADER_SIZE = 50
TNF_RECORD_SIZE = 360
SH_STOCK = re.compile(r"sh(?:600|601|603|605|688|689)\d{3}$")
SZ_STOCK = re.compile(r"sz(?:000|001|002|003|300|301)\d{3}$")


@dataclass(frozen=True)
class DayFileMetadata:
    source_code: str
    path: Path
    file_size: int
    record_count: int
    first_date: str | None
    last_date: str | None
    latest_close: float | None
    avg_amount_20: float | None
    latest_volume: int | None
    record_size_valid: bool
    latest_ohlc_valid: bool


def is_stock_code(source_code: str) -> bool:
    return bool(SH_STOCK.fullmatch(source_code) or SZ_STOCK.fullmatch(source_code))


def canonical_id(source_code: str) -> str:
    return f"{source_code[:2].upper()}.{source_code[2:]}"


def _date_text(value: int) -> str | None:
    try:
        return datetime.strptime(str(value), "%Y%m%d").date().isoformat()
    except (TypeError, ValueError):
        return None


def read_day_file_metadata(path: Path) -> DayFileMetadata:
    size = path.stat().st_size
    valid_size = size > 0 and size % DAY_RECORD_SIZE == 0
    count = size // DAY_RECORD_SIZE if valid_size else 0
    if count == 0:
        return DayFileMetadata(path.stem, path, size, 0, None, None, None, None, None, valid_size, False)

    with path.open("rb") as handle:
        first_raw = handle.read(DAY_RECORD_SIZE)
        tail_count = min(20, count)
        handle.seek(-tail_count * DAY_RECORD_SIZE, 2)
        tail_raw = handle.read(tail_count * DAY_RECORD_SIZE)

    first = DAY_STRUCT.unpack(first_raw)
    tail = [DAY_STRUCT.unpack_from(tail_raw, i * DAY_RECORD_SIZE) for i in range(tail_count)]
    last = tail[-1]
    amounts = np.array([row[5] for row in tail], dtype=float)
    amounts = amounts[np.isfinite(amounts) & (amounts >= 0)]
    o, h, low, close = (last[1] / 100.0, last[2] / 100.0, last[3] / 100.0, last[4] / 100.0)
    ohlc_valid = min(o, h, low, close) > 0 and h >= max(o, close, low) and low <= min(o, close, h)
    return DayFileMetadata(
        source_code=path.stem,
        path=path,
        file_size=size,
        record_count=count,
        first_date=_date_text(first[0]),
        last_date=_date_text(last[0]),
        latest_close=close,
        avg_amount_20=float(amounts.mean()) if amounts.size else None,
        latest_volume=int(last[6]),
        record_size_valid=valid_size,
        latest_ohlc_valid=ohlc_valid,
    )


def scan_stock_day_files(tdx_root: Path) -> list[DayFileMetadata]:
    results: list[DayFileMetadata] = []
    for market in ("sh", "sz"):
        folder = tdx_root / "vipdoc" / market / "lday"
        if not folder.exists():
            continue
        for path in sorted(folder.glob("*.day")):
            if is_stock_code(path.stem):
                results.append(read_day_file_metadata(path))
    return results


def read_day_dates(path: Path) -> np.ndarray:
    size = path.stat().st_size
    if size == 0 or size % DAY_RECORD_SIZE:
        raise ValueError(f"无效的通达信日线文件：{path}")
    raw = np.fromfile(path, dtype=np.dtype([
        ("date", "<u4"), ("open", "<u4"), ("high", "<u4"), ("low", "<u4"),
        ("close", "<u4"), ("amount", "<f4"), ("volume", "<u4"), ("reserved", "<u4"),
    ]))
    parsed = np.array([np.datetime64(_date_text(int(value)), "D") for value in raw["date"] if _date_text(int(value))])
    return np.unique(parsed)


def load_industry_names(tdx_root: Path) -> dict[str, str]:
    path = tdx_root / "incon.dat"
    result: dict[str, str] = {}
    for line in path.read_text(encoding="gb18030", errors="replace").splitlines():
        if "|" not in line:
            continue
        code, name = line.split("|", 1)
        result[code.strip()] = name.strip()
    return result


def load_tdx_security_names(tdx_root: Path) -> dict[str, str]:
    """Read the current SH/SZ security names from TDX market tables."""
    result: dict[str, str] = {}
    folder = tdx_root / "T0002" / "hq_cache"
    for market, filename in (("sh", "shs.tnf"), ("sz", "szs.tnf")):
        path = folder / filename
        if not path.exists():
            continue
        raw = path.read_bytes()
        payload = raw[TNF_HEADER_SIZE:]
        if len(payload) % TNF_RECORD_SIZE != 0:
            continue
        for offset in range(0, len(payload), TNF_RECORD_SIZE):
            row = payload[offset:offset + TNF_RECORD_SIZE]
            code = row[:6].decode("ascii", errors="ignore")
            source_code = market + code
            if not is_stock_code(source_code):
                continue
            name = row[31:39].split(b"\x00", 1)[0].decode("gb18030", errors="replace").strip()
            if name:
                result[source_code] = name
    return result


def load_tdx_research_industry(tdx_root: Path) -> dict[str, dict[str, str | None]]:
    names = load_industry_names(tdx_root)
    cfg = tdx_root / "T0002" / "hq_cache" / "tdxhy.cfg"
    result: dict[str, dict[str, str | None]] = {}
    for line in cfg.read_text(encoding="gb18030", errors="replace").splitlines():
        parts = line.split("|")
        if len(parts) < 6 or parts[0] not in {"0", "1"}:
            continue
        raw_code = parts[1].strip()
        industry = parts[5].strip()
        if not re.fullmatch(r"\d{6}", raw_code) or not industry.startswith("X"):
            continue
        source_code = ("sh" if parts[0] == "1" else "sz") + raw_code
        l1 = industry[:3] if len(industry) >= 3 else industry
        l2 = industry[:5] if len(industry) >= 5 else industry
        result[source_code] = {
            "industry_l1_code": l1,
            "industry_l1_name": names.get(l1),
            "industry_l2_code": l2,
            "industry_l2_name": names.get(l2),
            "industry_l3_code": industry,
            "industry_l3_name": names.get(industry),
        }
    return result