from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional


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

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


@dataclass
class SBDDJob:
    job_id: str
    job_name: str
    status: str
    protein_file_name: Optional[str] = None
    total_molecules: int = 0
    generated: 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) -> SBDDJob:
        return cls(
            job_id=data.get("job_id", ""),
            job_name=data.get("job_name", ""),
            status=data.get("status", ""),
            protein_file_name=data.get("protein_file_name"),
            total_molecules=data.get("total_molecules", 0),
            generated=data.get("generated", 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 SBDDJobStatus:
    success: bool
    message: Optional[str] = None
    job_id: Optional[str] = None
    denovo_job_id: Optional[str] = None
    job_name: Optional[str] = None
    protein_file_name: Optional[str] = None
    job_status: Optional[str] = None
    total_requested: int = 0
    total_generated: int = 0
    completed: int = 0
    failed: int = 0
    pending: int = 0
    pocket_coord: Optional[list[float]] = None
    radius: Optional[int] = None

    @classmethod
    def from_dict(cls, data: dict) -> SBDDJobStatus:
        return cls(
            success=data.get("success", False),
            message=data.get("message"),
            job_id=data.get("job_id"),
            denovo_job_id=data.get("denovo_job_id"),
            job_name=data.get("job_name"),
            protein_file_name=data.get("protein_file_name"),
            job_status=data.get("job_status"),
            total_requested=data.get("total_requested", 0),
            total_generated=data.get("total_generated", 0),
            completed=data.get("completed", 0),
            failed=data.get("failed", 0),
            pending=data.get("pending", 0),
            pocket_coord=data.get("pocket_coord"),
            radius=data.get("radius"),
        )


@dataclass
class SBDDMolecule:
    id: str
    smiles: Optional[str] = None
    qed: Optional[float] = None
    logp: Optional[float] = None
    lipinski: Optional[int] = None
    validity: Optional[bool] = None
    connectivity: Optional[bool] = None
    vina_score: Optional[float] = None
    sdf_path: Optional[str] = None
    json_path: Optional[str] = None
    num_times_edited: Optional[int] = None
    is_mds_optimized: Optional[bool] = None
    alias: Optional[str] = None
    created_at: Optional[str] = None

    @classmethod
    def from_dict(cls, data: dict) -> SBDDMolecule:
        return cls(
            id=data.get("id", ""),
            smiles=data.get("smiles"),
            qed=data.get("qed"),
            logp=data.get("logp"),
            lipinski=data.get("lipinski"),
            validity=data.get("validity"),
            connectivity=data.get("connectivity"),
            vina_score=data.get("vina_score"),
            sdf_path=data.get("sdf_path"),
            json_path=data.get("json_path"),
            num_times_edited=data.get("num_times_edited"),
            is_mds_optimized=data.get("is_mds_optimized"),
            alias=data.get("alias"),
            created_at=data.get("created_at"),
        )


@dataclass
class SBDDResults:
    success: bool
    message: Optional[str] = None
    job_name: Optional[str] = None
    protein_file_name: Optional[str] = None
    total: int = 0
    limit: int = 50
    offset: int = 0
    molecules: list[SBDDMolecule] = field(default_factory=list)

    @classmethod
    def from_dict(cls, data: dict) -> SBDDResults:
        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"),
            total=data.get("total", 0),
            limit=data.get("limit", 50),
            offset=data.get("offset", 0),
            molecules=[SBDDMolecule.from_dict(m) for m in data.get("molecules", [])],
        )


@dataclass
class SBDDMoleculeDetail:
    success: bool
    message: Optional[str] = None
    id: Optional[str] = None
    smiles: Optional[str] = None
    sdf_content: Optional[str] = None
    json_result: Optional[dict] = None
    qed: Optional[float] = None
    logp: Optional[float] = None
    lipinski: Optional[int] = None
    validity: Optional[bool] = None
    connectivity: Optional[bool] = None
    vina_score: Optional[float] = None
    sdf_path: Optional[str] = None
    json_path: Optional[str] = None
    num_times_edited: Optional[int] = None
    is_mds_optimized: Optional[bool] = None
    alias: Optional[str] = None
    created_at: Optional[str] = None
    updated_at: Optional[str] = None

    @classmethod
    def from_dict(cls, data: dict) -> SBDDMoleculeDetail:
        return cls(
            success=data.get("success", False),
            message=data.get("message"),
            id=data.get("id"),
            smiles=data.get("smiles"),
            sdf_content=data.get("sdf_content"),
            json_result=data.get("json_result"),
            qed=data.get("qed"),
            logp=data.get("logp"),
            lipinski=data.get("lipinski"),
            validity=data.get("validity"),
            connectivity=data.get("connectivity"),
            vina_score=data.get("vina_score"),
            sdf_path=data.get("sdf_path"),
            json_path=data.get("json_path"),
            num_times_edited=data.get("num_times_edited"),
            is_mds_optimized=data.get("is_mds_optimized"),
            alias=data.get("alias"),
            created_at=data.get("created_at"),
            updated_at=data.get("updated_at"),
        )


@dataclass
class SBDDStatBlock:
    best: Optional[float] = None
    worst: Optional[float] = None
    avg: Optional[float] = None
    count: Optional[int] = None

    @classmethod
    def from_dict(cls, data: dict) -> SBDDStatBlock:
        return cls(
            best=data.get("best"),
            worst=data.get("worst"),
            avg=data.get("avg"),
            count=data.get("count"),
        )


@dataclass
class SBDDSummaryStats:
    total_molecules: int = 0
    vina_stats: Optional[SBDDStatBlock] = None
    qed_stats: Optional[SBDDStatBlock] = None
    logp_stats: Optional[dict] = None
    validity_rate: Optional[float] = None
    connectivity_rate: Optional[float] = None

    @classmethod
    def from_dict(cls, data: dict) -> SBDDSummaryStats:
        vina_raw = data.get("vina_stats")
        qed_raw = data.get("qed_stats")
        return cls(
            total_molecules=data.get("total_molecules", 0),
            vina_stats=SBDDStatBlock.from_dict(vina_raw) if vina_raw else None,
            qed_stats=SBDDStatBlock.from_dict(qed_raw) if qed_raw else None,
            logp_stats=data.get("logp_stats"),
            validity_rate=data.get("validity_rate"),
            connectivity_rate=data.get("connectivity_rate"),
        )


@dataclass
class SBDDTopMolecule:
    id: str
    smiles: Optional[str] = None
    vina_score: Optional[float] = None
    qed: Optional[float] = None
    logp: Optional[float] = None

    @classmethod
    def from_dict(cls, data: dict) -> SBDDTopMolecule:
        return cls(
            id=data.get("id", ""),
            smiles=data.get("smiles"),
            vina_score=data.get("vina_score"),
            qed=data.get("qed"),
            logp=data.get("logp"),
        )


@dataclass
class SBDDSummary:
    success: bool
    message: Optional[str] = None
    job_name: Optional[str] = None
    protein_file_name: Optional[str] = None
    summary: Optional[SBDDSummaryStats] = None
    top_molecules: list[SBDDTopMolecule] = field(default_factory=list)

    @classmethod
    def from_dict(cls, data: dict) -> SBDDSummary:
        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=SBDDSummaryStats.from_dict(summary_raw) if summary_raw else None,
            top_molecules=[SBDDTopMolecule.from_dict(m) for m in data.get("top_molecules", [])],
        )
