"""Fetching 3D protein structures: experimental structures UniProt
cross-references (RCSB PDB) when available, otherwise the public AlphaFold
DB's predicted model."""

from __future__ import annotations

from dataclasses import dataclass

import requests

UNIPROT_ENTRY_URL = "https://rest.uniprot.org/uniprotkb/{accession}.json"
ALPHAFOLD_PREDICTION_URL = "https://alphafold.ebi.ac.uk/api/prediction/{accession}"
RCSB_PDB_URL = "https://files.rcsb.org/download/{pdb_id}.pdb"


@dataclass
class PdbCrossReference:
    pdb_id: str
    method: str | None
    resolution: str | None
    chains: str | None


@dataclass
class AlphaFoldPrediction:
    pdb_url: str
    mean_plddt: float | None
    model_created_date: str | None


def fetch_pdb_cross_references(accession: str, timeout: float = 10.0) -> list[PdbCrossReference]:
    """Experimental structures UniProt cross-references to RCSB PDB for
    `accession` (often empty — most proteins have no solved structure)."""
    url = UNIPROT_ENTRY_URL.format(accession=accession.strip())
    response = requests.get(url, timeout=timeout)
    response.raise_for_status()
    entry = response.json()

    references = []
    for xref in entry.get("uniProtKBCrossReferences", []):
        if xref.get("database") != "PDB":
            continue
        properties = {p["key"]: p["value"] for p in xref.get("properties", [])}
        references.append(
            PdbCrossReference(
                pdb_id=xref.get("id"),
                method=properties.get("Method"),
                resolution=properties.get("Resolution"),
                chains=properties.get("Chains"),
            )
        )
    return references


def fetch_alphafold_prediction(accession: str, timeout: float = 10.0) -> AlphaFoldPrediction | None:
    """Metadata for `accession`'s AlphaFold DB prediction (including the
    *current* model file URLs — AlphaFold DB periodically republishes
    models under a new version number, so this is fetched from the API
    rather than guessed/hardcoded), or None if it has no model."""
    url = ALPHAFOLD_PREDICTION_URL.format(accession=accession.strip())
    response = requests.get(url, timeout=timeout)
    if response.status_code == 404:
        return None
    response.raise_for_status()
    entries = response.json()
    if not entries:
        return None
    entry = entries[0]
    return AlphaFoldPrediction(
        pdb_url=entry["pdbUrl"],
        mean_plddt=entry.get("globalMetricValue"),
        model_created_date=entry.get("modelCreatedDate"),
    )


def fetch_structure_text(url: str, timeout: float = 20.0) -> str:
    """Raw PDB-format text from any public URL (AlphaFold DB or RCSB)."""
    response = requests.get(url, timeout=timeout)
    response.raise_for_status()
    return response.text


def fetch_pdb_structure(pdb_id: str, timeout: float = 20.0) -> str:
    """An experimental structure from RCSB PDB, as PDB text."""
    return fetch_structure_text(RCSB_PDB_URL.format(pdb_id=pdb_id.strip()), timeout=timeout)