from __future__ import annotations

import re
from collections.abc import Mapping
from datetime import datetime
from io import BytesIO
from pathlib import Path
from typing import TYPE_CHECKING, Any, List, Union

from astropy.io import fits
from astropy.table import Row, Table

from .models import ScanRecord, ScanType, normalize_scan_type

if TYPE_CHECKING:
    from .client import Grad300Client


TABLE_BASE_COLUMNS = [
    "id",
    "scan_type",
    "file_name",
    "file_path",
    "size",
    "source",
    "ra",
    "dec",
    "date",
    "date_obs",
    "date_end",
    "gain",
    "comment",
]

DownloadResult = Union[Path, Table, fits.HDUList]


class ScansService:
    """High-level helper methods around the /scans API routes."""

    def __init__(self, client: "Grad300Client") -> None:
        self._client = client

    def list(
        self,
        scan_type: ScanType | str,
        *,
        offset: int = 0,
        limit: int = 100,
        source: str | None = None,
        scan_name: str | None = None,
        date_from: datetime | None = None,
        date_to: datetime | None = None,
    ) -> Table:
        records = self.list_records(
            scan_type,
            offset=offset,
            limit=limit,
            source=source,
            scan_name=scan_name,
            date_from=date_from,
            date_to=date_to,
        )
        return self._records_to_table(records)

    def list_records(
        self,
        scan_type: ScanType | str,
        *,
        offset: int = 0,
        limit: int = 100,
        source: str | None = None,
        scan_name: str | None = None,
        date_from: datetime | None = None,
        date_to: datetime | None = None,
    ) -> List[ScanRecord]:
        if limit <= 0:
            raise ValueError("'limit' must be greater than 0")
        if offset < 0:
            raise ValueError("'offset' must be greater than or equal to 0")

        scan_type_value = normalize_scan_type(scan_type)

        params: dict[str, Any] = {"offset": offset, "limit": limit}
        if source:
            params["source"] = source
        if scan_name:
            params["scan_name"] = scan_name
        if date_from:
            params["date_from"] = date_from.isoformat()
        if date_to:
            params["date_to"] = date_to.isoformat()

        payload = self._client.get_json(
            f"/scans/{scan_type_value.value}/", params=params
        )
        if not isinstance(payload, list):
            raise ValueError("Unexpected payload from scan listing endpoint")

        records: List[ScanRecord] = []
        for item in payload:
            if not isinstance(item, dict):
                raise ValueError("Unexpected scan item format from backend API")
            records.append(ScanRecord.from_api(scan_type_value, item))

        return records

    def query(
        self,
        *,
        scan_type: ScanType | str | None = None,
        source: str | None = None,
        scan_name: str | None = None,
        date_from: datetime | None = None,
        date_to: datetime | None = None,
        page_size: int = 100,
        max_results: int | None = None,
    ) -> Table:
        records = self.query_records(
            scan_type=scan_type,
            source=source,
            scan_name=scan_name,
            date_from=date_from,
            date_to=date_to,
            page_size=page_size,
            max_results=max_results,
        )
        return self._records_to_table(records)

    def query_records(
        self,
        *,
        scan_type: ScanType | str | None = None,
        source: str | None = None,
        scan_name: str | None = None,
        date_from: datetime | None = None,
        date_to: datetime | None = None,
        page_size: int = 100,
        max_results: int | None = None,
    ) -> List[ScanRecord]:
        if page_size <= 0:
            raise ValueError("'page_size' must be greater than 0")
        if max_results is not None and max_results <= 0:
            raise ValueError("'max_results' must be greater than 0 when provided")
        if date_from and date_to and date_to < date_from:
            raise ValueError("'date_to' must be greater than or equal to 'date_from'")

        scan_types = [normalize_scan_type(scan_type)] if scan_type else list(ScanType)
        results: List[ScanRecord] = []

        for current_scan_type in scan_types:
            offset = 0
            while True:
                batch = self.list_records(
                    current_scan_type,
                    offset=offset,
                    limit=page_size,
                    source=source,
                    scan_name=scan_name,
                    date_from=date_from,
                    date_to=date_to,
                )

                if not batch:
                    break

                for record in batch:
                    if self._matches_local_filters(
                        record,
                        source=source,
                        scan_name=scan_name,
                        date_from=date_from,
                        date_to=date_to,
                    ):
                        results.append(record)
                        if max_results is not None and len(results) >= max_results:
                            return results

                if len(batch) < page_size:
                    break
                offset += len(batch)

        return results

    def query_source(
        self,
        source: str,
        *,
        scan_type: ScanType | str | None = None,
        page_size: int = 100,
        max_results: int | None = None,
    ) -> Table:
        return self.query(
            scan_type=scan_type,
            source=source,
            page_size=page_size,
            max_results=max_results,
        )

    def query_scan_name(
        self,
        scan_name: str,
        *,
        scan_type: ScanType | str | None = None,
        page_size: int = 100,
        max_results: int | None = None,
    ) -> Table:
        return self.query(
            scan_type=scan_type,
            scan_name=scan_name,
            page_size=page_size,
            max_results=max_results,
        )

    def query_date_range(
        self,
        *,
        date_from: datetime | None = None,
        date_to: datetime | None = None,
        scan_type: ScanType | str | None = None,
        page_size: int = 100,
        max_results: int | None = None,
    ) -> Table:
        return self.query(
            scan_type=scan_type,
            date_from=date_from,
            date_to=date_to,
            page_size=page_size,
            max_results=max_results,
        )

    def download(
        self,
        record: ScanRecord | Mapping[str, Any] | Row,
        *,
        destination_dir: str | Path | None = None,
        overwrite: bool = False,
    ) -> DownloadResult:
        scan_type, scan_id, file_name = self._extract_download_fields(record)
        return self.download_by_id(
            scan_type=scan_type,
            scan_id=scan_id,
            destination_dir=destination_dir,
            file_name=file_name,
            overwrite=overwrite,
        )

    def download_by_id(
        self,
        *,
        scan_type: ScanType | str,
        scan_id: int,
        destination_dir: str | Path | None = None,
        file_name: str | None = None,
        overwrite: bool = False,
    ) -> DownloadResult:
        scan_type_value = normalize_scan_type(scan_type)
        response = self._client.request(
            "GET",
            f"/scans/{scan_type_value.value}/{scan_id}/download",
        )

        if destination_dir is None:
            return _parse_download_content(scan_type_value, response.content)

        output_dir = Path(destination_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        resolved_name = file_name or _filename_from_disposition(
            response.headers.get("content-disposition")
        )
        if not resolved_name:
            resolved_name = f"{scan_type_value.value}_{scan_id}.fits"

        output_path = output_dir / resolved_name
        if output_path.exists() and not overwrite:
            raise FileExistsError(
                f"Destination file already exists: {output_path}. Set overwrite=True to replace it."
            )

        output_path.write_bytes(response.content)
        return output_path

    def download_by_scan_name(
        self,
        scan_name: str,
        *,
        scan_type: ScanType | str | None = None,
        destination_dir: str | Path | None = None,
        overwrite: bool = False,
    ) -> DownloadResult:
        candidates = self.query(
            scan_type=scan_type,
            scan_name=scan_name,
        )
        exact_matches = [
            row for row in candidates if str(row["file_name"]) == scan_name
        ]

        if not exact_matches:
            raise LookupError(f"No scan found with exact name '{scan_name}'")
        if len(exact_matches) > 1:
            raise LookupError(
                f"Multiple scans found with name '{scan_name}'. Narrow down with scan_type/date filters."
            )

        return self.download(
            exact_matches[0],
            destination_dir=destination_dir,
            overwrite=overwrite,
        )

    @classmethod
    def _records_to_table(cls, records: List[ScanRecord]) -> Table:
        if not records:
            return Table(
                names=TABLE_BASE_COLUMNS,
                dtype=["O"] * len(TABLE_BASE_COLUMNS),
            )

        metadata_columns = sorted(
            {key for record in records for key in record.metadata.keys()}
        )
        column_names = TABLE_BASE_COLUMNS + metadata_columns

        rows: List[dict[str, Any]] = []
        for record in records:
            row: dict[str, Any] = {
                "id": record.id,
                "scan_type": record.scan_type.value,
                "file_name": record.file_name,
                "file_path": record.file_path,
                "size": record.size,
                "source": record.source,
                "ra": record.ra,
                "dec": record.dec,
                "date": record.date,
                "date_obs": record.date_obs,
                "date_end": record.date_end,
                "gain": record.gain,
                "comment": record.comment,
            }
            for key in metadata_columns:
                row[key] = record.metadata.get(key)
            rows.append(row)

        return Table(rows=rows, names=column_names)

    @staticmethod
    def _extract_download_fields(
        record: ScanRecord | Mapping[str, Any] | Row,
    ) -> tuple[ScanType | str, int, str | None]:
        if isinstance(record, ScanRecord):
            return record.scan_type, record.id, record.file_name

        scan_type = _record_get(record, "scan_type")
        scan_id = _record_get(record, "id")
        file_name = _record_get(record, "file_name")

        if scan_type is None:
            raise TypeError("Record must include a 'scan_type' field")
        if scan_id is None:
            raise TypeError("Record must include an 'id' field")

        parsed_file_name = str(file_name) if file_name is not None else None
        return scan_type, int(scan_id), parsed_file_name

    @staticmethod
    def _matches_local_filters(
        record: ScanRecord,
        *,
        source: str | None,
        scan_name: str | None,
        date_from: datetime | None,
        date_to: datetime | None,
    ) -> bool:
        if source and source.casefold() not in record.source.casefold():
            return False
        if scan_name and scan_name.casefold() not in record.file_name.casefold():
            return False
        if date_from and record.date < date_from:
            return False
        if date_to and record.date > date_to:
            return False
        return True


def _filename_from_disposition(content_disposition: str | None) -> str | None:
    if not content_disposition:
        return None

    match = re.search(r'filename="?([^";]+)"?', content_disposition)
    if match:
        return match.group(1)
    return None


def _record_get(record: Mapping[str, Any] | Row, key: str) -> Any:
    if isinstance(record, Mapping):
        return record.get(key)

    try:
        return record[key]
    except Exception:
        return None


def _parse_download_content(
    scan_type: ScanType,
    content: bytes,
) -> Table | fits.HDUList:
    if scan_type == ScanType.TPI:
        return _parse_tpi_table(content)

    return fits.open(BytesIO(content), memmap=False)


def _parse_tpi_table(content: bytes) -> Table:
    try:
        return Table.read(BytesIO(content), format="fits")
    except Exception:
        with fits.open(BytesIO(content), memmap=False) as hdul:
            for hdu in hdul:
                data = getattr(hdu, "data", None)
                if data is None:
                    continue
                try:
                    return Table(data)
                except Exception:
                    continue

    raise ValueError("Unable to parse TPI FITS content as an astropy Table")
