from pathlib import Path
import re
from enum import Enum
from typing import Annotated, Literal, Tuple, List, Optional

from pydantic import BaseModel, Field, field_validator

from aaredaqlib.coordinate import Coordinate
from aaredaqlib.diffraction_geometry import DiffractionGeometry
from aaredaqlib.sample_geometry import SampleGeometryModel

class StagePositionEnum(Enum):
    MEASURE = 0
    PARK = 1
    DOWN = 2
    UNKNOWN = 3

class TokenData(BaseModel):
    sub: str # Username
    pgroups: List[str]
    session: int
    staff: bool = False

class DewarAddress(BaseModel):
    segment: Literal["A", "B", "C", "D", "E", "F", "X", "R"]
    pos: Annotated[int, Field(ge=1, le=5)]


class SampleDewarAddress(BaseModel):
    puck: DewarAddress
    pin: Annotated[int, Field(ge=1, le=16)]


# From the database for puck loading
class PuckInfo(BaseModel):
    db_id: int
    puck_name: str
    dewar_name: str
    user: str = ""
    location: Optional[DewarAddress] = None


# From TELL to database after loading
class PuckLoadedInfo(BaseModel):
    puck_name: str
    location: DewarAddress

class DataCollectionParameters(BaseModel):
    directory: Optional[str] = None
    oscillation: Optional[float] = None  # Only accept positive float
    exposure: Optional[float] = None  # Only accept positive floats between 0 and 1
    totalrange: Optional[int] = None  # Only accept positive integers between 0 and 360
    transmission: Optional[
        int
    ] = None  # Only accept positive integers between 0 and 100
    targetresolution: Optional[float] = None  # Only accept positive float
    aperture: Optional[str] = None  # Optional string field
    datacollectiontype: Optional[
        str
    ] = None  # Only accept "standard", other types might be added later
    processingpipeline: Optional[
        str
    ] = ""  # Only accept "gopy", "autoproc", "xia2dials"
    spacegroupnumber: Optional[
        int
    ] = None  # Only accept positive integers between 1 and 230
    cellparameters: Optional[
        str
    ] = None  # Must be a set of six positive floats or integers
    rescutkey: Optional[str] = None  # Only accept "is" or "cchalf"
    rescutvalue: Optional[
        float
    ] = None  # Must be a positive float if rescutkey is provided
    userresolution: Optional[float] = None
    pdbid: Optional[
        str
    ] = ""  # Accepts either the format of the protein data bank code or {provided}
    autoprocfull: Optional[bool] = None
    procfull: Optional[bool] = None
    adpenabled: Optional[bool] = None
    noano: Optional[bool] = None
    ffcscampaign: Optional[bool] = None
    trustedhigh: Optional[float] = None  # Should be a float between 0 and 2.0
    autoprocextraparams: Optional[str] = None  # Optional string field
    chiphiangles: Optional[float] = None  # Optional float field between 0 and 30
    dose: Optional[float] = None  # Optional float field
    cloud: bool = True
    pdbmodel: Optional[str] = None

    def to_dict(self):
        """Convert the model instance to a dictionary."""
        return self.dict(
            exclude_unset=True
        )  # Use this built-in method for serialization

    class Config:
        from_attributes = True

    @field_validator("directory", mode="after")
    @classmethod
    def directory_characters(cls, v):

        # Default directory value if empty
        if not v:  # Handles None or empty cases
            default_value = "{sgPuck}/{sgPosition}"
            return default_value

        # Strip trailing slashes and store original value for comparison
        v = str(v).strip("/")  # Ensure it's a string and no trailing slashes
        original_value = v

        # Replace spaces with underscores
        v = v.replace(" ", "_")

        # Validate directory pattern with macros and allowed characters
        valid_macros = [
            "{date}",
            "{prefix}",
            "{sgPuck}",
            "{sgPosition}",
            "{beamline}",
            "{sgPrefix}",
            "{sgPriority}",
            "{protein}",
            "{method}",
        ]
        valid_macro_pattern = re.compile(
            "|".join(re.escape(macro) for macro in valid_macros)
        )

        # Check if the value contains valid macros
        allowed_chars_pattern = "[a-z0-9_.+-/]"
        v_without_macros = valid_macro_pattern.sub("macro", v)

        allowed_path_pattern = re.compile(
            f"^(({allowed_chars_pattern}+|macro)*/*)*$", re.IGNORECASE
        )
        if not allowed_path_pattern.match(v_without_macros):
            raise ValueError(
                f"'{v}' is not valid. Value must be a valid path or macro."
            )
        return v

    @field_validator("aperture", mode="before")
    @classmethod
    def aperture_selection(cls, v):
        if v is not None:
            try:
                v = int(float(v))
                if v not in {1, 2, 3}:
                    raise ValueError(f" '{v}' is not valid. Value must be 1, 2, or 3.")
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid. Value must be 1, 2, or 3."
                ) from e
        return v

    @field_validator("oscillation", mode="before")
    @classmethod
    def positive_float_validator(cls, v):
        if v is None:
            return None
        try:
            v = float(v)
            if v <= 0:
                raise ValueError(f"'{v}' is not valid. Value must be a positive float.")
        except (ValueError, TypeError) as e:
            raise ValueError(
                f"'{v}' is not valid. Value must be a positive float."
            ) from e
        return v

    @field_validator("exposure", mode="before")
    @classmethod
    def exposure_in_range(cls, v):
        if v is not None:
            try:
                v = float(v)
                if not (0 <= v <= 1):
                    raise ValueError(
                        f" '{v}' is not valid. Value must be a float between 0 and 1."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid. Value must be a float between 0 and 1."
                ) from e
        return v

    @field_validator("totalrange", mode="before")
    @classmethod
    def totalrange_in_range(cls, v):
        if v is not None:
            try:
                v = int(v)
                if not (0 <= v <= 360):
                    raise ValueError(
                        f" '{v}' is not valid."
                        f"Value must be an integer between 0 and 360."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid."
                    f"Value must be an integer between 0 and 360."
                ) from e
        return v

    @field_validator("transmission", mode="before")
    @classmethod
    def transmission_fraction(cls, v):
        if v is not None:
            try:
                v = int(v)
                if not (0 <= v <= 100):
                    raise ValueError(
                        f" '{v}' is not valid."
                        f"Value must be an integer between 0 and 100."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid."
                    f"Value must be an integer between 0 and 100."
                ) from e
        return v

    @field_validator("datacollectiontype", mode="before")
    @classmethod
    def datacollectiontype_allowed(cls, v):
        allowed = {"standard"}  # Other types of data collection might be added later
        if v and v.lower() not in allowed:
            raise ValueError(f" '{v}' is not valid." f"Value must be one of {allowed}.")
        return v

    @field_validator("processingpipeline", mode="before")
    @classmethod
    def processingpipeline_allowed(cls, v):
        allowed = {"aareproc", "autoproc"}
        if v and v.lower() not in allowed:
            raise ValueError(f" '{v}' is not valid." f"Value must be one of {allowed}.")
        return v

    @field_validator("spacegroupnumber", mode="before")
    @classmethod
    def spacegroupnumber_allowed(cls, v):
        if v is not None:
            try:
                v = int(v)
                if not (1 <= v <= 230):
                    raise ValueError(
                        f" '{v}' is not valid."
                        f"Value must be an integer between 1 and 230."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid."
                    f"Value must be an integer between 1 and 230."
                ) from e
        return v

    @field_validator("cellparameters", mode="before")
    @classmethod
    def cellparameters_format(cls, v):
        if v:
            # Replace commas with spaces, then split on whitespace
            tokens = v.replace(",", " ").split()
            try:
                values = [float(i) for i in tokens]
            except ValueError:
                raise ValueError(
                    f" '{v}' is not valid."
                    " Value must be a set of six positive floats"
                    " or integers (separated by space or comma)."
                )
            if len(values) != 6 or any(val <= 0 for val in values):
                raise ValueError(
                    f" '{v}' is not valid."
                    " Value must be a set of six positive floats"
                    " or integers (separated by space or comma)."
                )
        return v

    # @field_validator("rescutkey", "rescutvalue", mode="before")
    # @classmethod
    # def rescutkey_value_pair(cls, values):
    #    rescutkey = values.get("rescutkey")
    #    rescutvalue = values.get("rescutvalue")
    #    if rescutkey and rescutvalue:
    #        if rescutkey not in {"is", "cchalf"}:
    #            raise ValueError("Rescutkey must be either 'is' or 'cchalf'")
    #        if not isinstance(rescutvalue, float) or rescutvalue <= 0:
    #            raise ValueError(
    #                "Rescutvalue must be a positive float if rescutkey is provided"
    #            )
    #    return values

    @field_validator("trustedhigh", mode="before")
    @classmethod
    def trustedhigh_allowed(cls, v):
        if v is not None:
            try:
                v = float(v)
                if not (0 <= v <= 2.0):
                    raise ValueError(
                        f" '{v}' is not valid."
                        f"Value must be a float between 0 and 2.0."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid." f"Value must be a float between 0 and 2.0."
                ) from e
        return v

    @field_validator("chiphiangles", mode="before")
    @classmethod
    def chiphiangles_allowed(cls, v):
        if v is not None:
            try:
                v = float(v)
                if not (0 <= v <= 30):
                    raise ValueError(
                        f" '{v}' is not valid."
                        f"Value must be a float between 0 and 30."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid. Value must be a float between 0 and 30."
                ) from e
        return v

    @field_validator("dose", mode="before")
    @classmethod
    def dose_positive(cls, v):
        if v is not None:
            try:
                v = float(v)
                if v <= 0:
                    raise ValueError(
                        f" '{v}' is not valid. Value must be a positive float."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid. Value must be a positive float."
                ) from e
        return v

    @field_validator("pdbmodel", mode="after")
    @classmethod
    def validate_filepath(cls, v):
        if v is None:
            return v
        v_str = str(v)  # Ensure v is a string for further checks
        if any(c in v_str for c in '<>:"|?*'):
            raise ValueError("File path contains invalid characters.")
        path = Path(v_str)
        if not path.parts:
            raise ValueError("Not a valid path.")
        return v_str  # Return as string for JSON serialization

    @field_validator("cloud", mode="before")
    @classmethod
    def coerce_cloud_default(cls, v):
        if v in ("", None):
            return True
        if isinstance(v, bool):
            return v
        v_str = str(v).strip().lower()
        if v_str in {"true", "yes", "1"}:
            return True
        if v_str in {"false", "no", "0"}:
            return False
        raise ValueError("cloud must be blank for default, or True/False")

# From database to TELL after loading
class SampleShortInfo(BaseModel):
    db_id: int
    puck_name: str
    dewar_name: str
    sample_name: str
    aaredb_params: Optional[DataCollectionParameters] = None
    user: str = ""
    pin: Annotated[int, Field(ge=1, le=16)]
    location: DewarAddress | None = None
    priority: Optional[float] = 1.0
    comment: Optional[str] = None
    mount_count: int = 0
    def tell_address(self) -> SampleDewarAddress:
        return SampleDewarAddress(puck=self.location, pin=self.pin)

    def loc_str(self) -> str:
        if self.location is None:
            return "-"
        else:
            return f"{self.location.segment}{self.location.pos}-{self.pin}"

    def loc_str_sort(self) -> str:
        if self.location is None:
            return ""
        else:
            return f"{self.location.segment}{self.location.pos}-{self.pin:02d}"


class SampleShortInfoList(BaseModel):
    s: List[SampleShortInfo]

#class ReferencePuckInfo(BaseModel):
#    test_sample = SampleShortInfo(db_id=-1, puck_name="test_puck", dewar_name="test_dewar", sample_name="test_sample",
#                                  pin=11, location=DewarAddress(segment="X", pos=1))

class BeamMarkCoeffModel(BaseModel):
    """
    Model to calculate beam center for a given zoom level based on a quadratic approximation
    For both x and y it contains tuple of coefficients a, b, c (a*zoom^2 + b*zoom + c)
    Default is beam center in 1000,1000 at any zoom level
    """

    coeff_x: Tuple[float, float, float] = (0, 0, 1000.0)
    coeff_y: Tuple[float, float, float] = (0, 0, 1000.0)

    def apply(self, zoom: float):
        return Coordinate(
            x=self.coeff_x[0] * zoom**2 + self.coeff_x[1] * zoom + self.coeff_x[2],
            y=self.coeff_y[0] * zoom**2 + self.coeff_y[1] * zoom + self.coeff_y[2],
        )

class BoundingBoxModel(BaseModel):
    top_x: float
    top_y: float
    bottom_x: float
    bottom_y: float


class LoopCenteringZoomModelElem(BaseModel):
    zoom_value: float
    sam_cam_gain: float
    sam_cam_exp: float


class LoopCenteringZoomModel(BaseModel):
        z: List[LoopCenteringZoomModelElem]

# zoom gain exp (X06DA)
#    1  50 0.05
#  200  50 0.05
#  500  70 0.05
#  700 150 0.05
#  800 100 0.10
# 1000 200 0.10


def def_loop_centering_zoom() -> LoopCenteringZoomModel:
    return LoopCenteringZoomModel(
        z=[
            LoopCenteringZoomModelElem(
                zoom_value=1, sam_cam_gain=0, sam_cam_exp=0.05,
            )
        ]
    )


class BeamlineStateEnum(Enum):
    Maintenance = 1
    SampleExchange = 2
    SampleAlignment = 3
    DataCollection = 4
    DewarTransfer = 5
    XrayFluorescence = 6
    BeamLocation = 7
    Moving = 8
    RobotSampleExchange = 9
    # Busy state is used whenever transition between two states happen or mounting/data collection procedure happens


class SessionsStateEnum(Enum):
    OwnedByYou = 1
    OwnedByElse = 2
    Vacant = 3

class SampleCameraSettings(BaseModel):
    gain: float
    exposure: float


class ZoomModel(BaseModel):
    z:dict[float, SampleCameraSettings]

    def get_camera_settings(self, zoom_value: float) -> SampleCameraSettings | tuple[float, float]:
        if not self.z:
            raise ValueError("No zoom data available")

        if zoom_value in self.z:
            elem = self.z[zoom_value]
            return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)

        sorted_zooms = sorted(self.z.keys())

        if zoom_value <= sorted_zooms[0]:
            elem = self.z[sorted_zooms[0]]
            return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)

        if zoom_value >= sorted_zooms[-1]:
            elem = self.z[sorted_zooms[-1]]
            return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)
        print('interpolating zoom')
        for i in range(len(sorted_zooms) - 1):
            if sorted_zooms[i] <= zoom_value <= sorted_zooms[i + 1]:
                lower_zoom = sorted_zooms[i]
                upper_zoom = sorted_zooms[i + 1]

                lower_elem = self.z[lower_zoom]
                upper_elem = self.z[upper_zoom]

                t = (zoom_value - lower_zoom) / (upper_zoom - lower_zoom)

                interpolated_gain = lower_elem.gain + t * (upper_elem.gain - lower_elem.gain)
                interpolated_exp = lower_elem.exposure + t * (upper_elem.exposure - lower_elem.exposure)

                return SampleCameraSettings(gain=interpolated_gain, exposure=interpolated_exp)

        closest = min(self.z.keys(), key=lambda x: abs(x - zoom_value))
        elem = self.z[closest]
        return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)

def def_zoom() -> ZoomModel:
    return ZoomModel(
        z={
            1: SampleCameraSettings(gain=0, exposure=0.05),
            280: SampleCameraSettings(gain=0, exposure=0.05),
            500: SampleCameraSettings(gain=0, exposure=0.05),
            700: SampleCameraSettings(gain=0, exposure=0.05),
            800: SampleCameraSettings(gain=0, exposure=0.10),
            1000: SampleCameraSettings(gain=0, exposure=0.20)
        }
    )

def def_bl_zoom() -> ZoomModel:
    return ZoomModel(
        z={
            1: SampleCameraSettings(gain=0, exposure=0.002),
            280: SampleCameraSettings(gain=0, exposure=0.002),
            500: SampleCameraSettings(gain=0, exposure=0.002),
            700: SampleCameraSettings(gain=0, exposure=0.002),
            800: SampleCameraSettings(gain=0, exposure=0.002),
            1000: SampleCameraSettings(gain=0, exposure=0.005)
        }
    )

class AutofocusSettings(BaseModel):
    center_x_pxl: float
    center_y_pxl: float
    radius_pxl: float
    z_range_um: float
    z_steps: int

class BeamlineStatus(BaseModel):
    name: str
    ring_current_mA: float
    light: Annotated[float, Field(ge=0.0, le=100.0)]
    cryojet_K: float
    shutter_open: bool
    flux_ph_s: float
    sample_camera: SampleCameraSettings
    transmission: Annotated[float, Field(ge=0.0, le=1.0)] | None
    zoom: float

class SessionStatus(BaseModel):
    session: SessionsStateEnum = SessionsStateEnum.Vacant
    current_pgroup: str | None = None
    staff: bool = False

class DAQStatusModel(BaseModel):
    geom: SampleGeometryModel
    diffraction: DiffractionGeometry
    bl: BeamlineStatus
    state: BeamlineStateEnum
    busy: bool
    sample: SampleShortInfo | None = None
    session: SessionStatus
    box: BoundingBoxModel | None = None

class BeamlineSettingsModel(BaseModel):
    dtz_max: float | None = 1600.0
    dtz_min: float | None = 120.0
    dtz_collection: float | None = 130.0
    dtz_park: float | None = 150.0
    dtz_wash_sample_distance: float | None = 150.0
    dtz_bsz_safety_margin: float | None = 50.0
    bsz: float | None = 25.0

    camera_max_magnification: float | None = 1.0
    camera_min_magnification: float | None = 500.0
    camera_translation_factor_a: float | None = 0.00253
    camera_translation_factor_b: float | None = 255.0

class CryojetSettingsModel(BaseModel):
    cryojet_park_position: float | None = 12.0
    cryojet_measurement_position: float | None = 5.0
    cryojet_in_use: bool | None = True