from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional


@dataclass
class PocketInfo:
    pocket_id: int
    center: list[float]
    score: float
    druggability_score: Optional[float] = None
    volume: Optional[float] = None
    alpha_spheres: Optional[int] = None
    total_sasa: Optional[float] = None
    polar_sasa: Optional[float] = None
    apolar_sasa: Optional[float] = None
    hydrophobicity_score: Optional[float] = None
    polarity_score: Optional[int] = None
    charge_score: Optional[int] = None
    pocket_name: Optional[str] = None

    @classmethod
    def from_dict(cls, data: dict) -> PocketInfo:
        return cls(
            pocket_id=data.get("pocket_id", 0),
            center=data.get("center", []),
            score=data.get("score", 0.0),
            druggability_score=data.get("druggability_score"),
            volume=data.get("volume"),
            alpha_spheres=data.get("alpha_spheres"),
            total_sasa=data.get("total_sasa"),
            polar_sasa=data.get("polar_sasa"),
            apolar_sasa=data.get("apolar_sasa"),
            hydrophobicity_score=data.get("hydrophobicity_score"),
            polarity_score=data.get("polarity_score"),
            charge_score=data.get("charge_score"),
            pocket_name=data.get("pocket_name"),
        )


@dataclass
class PocketDetectionResult:
    success: bool
    message: Optional[str] = None
    num_pockets: int = 0
    pockets: list[PocketInfo] = field(default_factory=list)

    @classmethod
    def from_dict(cls, data: dict) -> PocketDetectionResult:
        return cls(
            success=data.get("success", False),
            message=data.get("message"),
            num_pockets=data.get("num_pockets", 0),
            pockets=[PocketInfo.from_dict(p) for p in data.get("pockets", [])],
        )


@dataclass
class DockingJobSubmission:
    success: bool
    message: str
    job_id: Optional[str] = None

    @classmethod
    def from_dict(cls, data: dict) -> DockingJobSubmission:
        return cls(
            success=data.get("success", False),
            message=data.get("message", ""),
            job_id=data.get("job_id"),
        )


@dataclass
class DockingJob:
    job_id: str
    job_name: str
    status: str
    total_ligands: int = 0
    completed: int = 0
    failed: int = 0
    pending: int = 0
    created_at: Optional[str] = None
    completed_at: Optional[str] = None

    @classmethod
    def from_dict(cls, data: dict) -> DockingJob:
        return cls(
            job_id=data.get("job_id", ""),
            job_name=data.get("job_name", ""),
            status=data.get("status", ""),
            total_ligands=data.get("total_ligands", 0),
            completed=data.get("completed", 0),
            failed=data.get("failed", 0),
            pending=data.get("pending", 0),
            created_at=data.get("created_at"),
            completed_at=data.get("completed_at"),
        )


@dataclass
class DockingLigandStatus:
    ligand_name: str
    status: str
    num_poses: Optional[int] = None
    best_score: Optional[float] = None
    pocket_coord: Optional[list[float]] = None
    box_sizes: Optional[list[float]] = None

    @classmethod
    def from_dict(cls, data: dict) -> DockingLigandStatus:
        return cls(
            ligand_name=data.get("ligand_name", ""),
            status=data.get("status", ""),
            num_poses=data.get("num_poses"),
            best_score=data.get("best_score"),
            pocket_coord=data.get("pocket_coord"),
            box_sizes=data.get("box_sizes"),
        )


@dataclass
class DockingJobStatus:
    success: bool
    message: Optional[str] = None
    job_id: Optional[str] = None
    job_name: Optional[str] = None
    protein_file_name: Optional[str] = None
    job_status: Optional[str] = None
    total_ligands: int = 0
    completed: int = 0
    failed: int = 0
    pending: int = 0
    ligands: list[DockingLigandStatus] = field(default_factory=list)

    @classmethod
    def from_dict(cls, data: dict) -> DockingJobStatus:
        return cls(
            success=data.get("success", False),
            message=data.get("message"),
            job_id=data.get("job_id"),
            job_name=data.get("job_name"),
            protein_file_name=data.get("protein_file_name"),
            job_status=data.get("job_status"),
            total_ligands=data.get("total_ligands", 0),
            completed=data.get("completed", 0),
            failed=data.get("failed", 0),
            pending=data.get("pending", 0),
            ligands=[DockingLigandStatus.from_dict(l) for l in data.get("ligands", [])],
        )


@dataclass
class InteractionBlock:
    hydrophobic_interactions: list[dict] = field(default_factory=list)
    hydrogen_bonds: list[dict] = field(default_factory=list)
    pi_stacking: list[dict] = field(default_factory=list)
    pi_cation: list[dict] = field(default_factory=list)
    salt_bridges: list[dict] = field(default_factory=list)
    water_bridges: list[dict] = field(default_factory=list)
    halogen_bonds: list[dict] = field(default_factory=list)
    metal_complexes: list[dict] = field(default_factory=list)

    @classmethod
    def from_dict(cls, data: dict) -> InteractionBlock:
        return cls(
            hydrophobic_interactions=data.get("hydrophobic_interactions", []),
            hydrogen_bonds=data.get("hydrogen_bonds", []),
            pi_stacking=data.get("pi_stacking", []),
            pi_cation=data.get("pi_cation", []),
            salt_bridges=data.get("salt_bridges", []),
            water_bridges=data.get("water_bridges", []),
            halogen_bonds=data.get("halogen_bonds", []),
            metal_complexes=data.get("metal_complexes", []),
        )


@dataclass
class PoseInfo:
    conf_id: int
    score: float
    mol_name: str = ""
    interactions: dict[str, InteractionBlock] = field(default_factory=dict)

    @classmethod
    def from_dict(cls, data: dict) -> PoseInfo:
        raw_interactions = data.get("interactions", {})
        interactions = {
            k: InteractionBlock.from_dict(v) for k, v in raw_interactions.items()
        }
        return cls(
            conf_id=data.get("conf_id", 0),
            score=data.get("score", 0.0),
            mol_name=data.get("mol_name", ""),
            interactions=interactions,
        )


@dataclass
class DockingResult:
    success: bool
    message: Optional[str] = None
    job_name: Optional[str] = None
    protein_file_name: Optional[str] = None
    ligand_name: Optional[str] = None
    status: Optional[str] = None
    best_score: Optional[float] = None
    num_poses: Optional[int] = None
    pocket_coord: Optional[list[float]] = None
    box_sizes: Optional[list[float]] = None
    sdf_content: Optional[str] = None
    pose_info: Optional[PoseInfo] = None

    @classmethod
    def from_dict(cls, data: dict) -> DockingResult:
        pose_raw = data.get("pose_info")
        return cls(
            success=data.get("success", False),
            message=data.get("message"),
            job_name=data.get("job_name"),
            protein_file_name=data.get("protein_file_name"),
            ligand_name=data.get("ligand_name"),
            status=data.get("status"),
            best_score=data.get("best_score"),
            num_poses=data.get("num_poses"),
            pocket_coord=data.get("pocket_coord"),
            box_sizes=data.get("box_sizes"),
            sdf_content=data.get("sdf_content"),
            pose_info=PoseInfo.from_dict(pose_raw) if pose_raw else None,
        )


@dataclass
class DockingResultSummaryItem:
    ligand_name: str
    best_score: Optional[float] = None
    num_poses: Optional[int] = None
    pocket_coord: Optional[list[float]] = None
    output_file_path: Optional[str] = None

    @classmethod
    def from_dict(cls, data: dict) -> DockingResultSummaryItem:
        return cls(
            ligand_name=data.get("ligand_name", ""),
            best_score=data.get("best_score"),
            num_poses=data.get("num_poses"),
            pocket_coord=data.get("pocket_coord"),
            output_file_path=data.get("output_file_path"),
        )


@dataclass
class DockingSummaryStats:
    total_ligands: int = 0
    best_score: Optional[float] = None
    worst_score: Optional[float] = None
    avg_score: Optional[float] = None

    @classmethod
    def from_dict(cls, data: dict) -> DockingSummaryStats:
        return cls(
            total_ligands=data.get("total_ligands", 0),
            best_score=data.get("best_score"),
            worst_score=data.get("worst_score"),
            avg_score=data.get("avg_score"),
        )


@dataclass
class DockingSummary:
    success: bool
    message: Optional[str] = None
    job_name: Optional[str] = None
    protein_file_name: Optional[str] = None
    summary: Optional[DockingSummaryStats] = None
    results: list[DockingResultSummaryItem] = field(default_factory=list)

    @classmethod
    def from_dict(cls, data: dict) -> DockingSummary:
        summary_raw = data.get("summary")
        return cls(
            success=data.get("success", False),
            message=data.get("message"),
            job_name=data.get("job_name"),
            protein_file_name=data.get("protein_file_name"),
            summary=DockingSummaryStats.from_dict(summary_raw) if summary_raw else None,
            results=[DockingResultSummaryItem.from_dict(r) for r in data.get("results", [])],
        )
