# genesis/molecule_viewer.py
import requests
from typing import Optional

def fetch_pdb(pdb_id: str) -> Optional[str]:
    """Fetch PDB file from RCSB."""
    url = f"https://files.rcsb.org/download/{pdb_id.upper()}.pdb"
    r = requests.get(url)
    if r.status_code == 200:
        return r.text
    return None

def fetch_alphafold(uniprot_id: str) -> Optional[str]:
    """Fetch PDB model from AlphaFold."""
    url = f"https://alphafold.ebi.ac.uk/files/AF-{uniprot_id}-F1-model_v4.pdb"
    r = requests.get(url)
    if r.status_code == 200:
        return r.text
    return None

def render_3dmol(pdb_data: str) -> str:
    """Generate HTML for 3Dmol.js viewer with biotech neon styling."""
    return f"""
    <div id="viewer" style="width: 100%; height: 600px; background-color: black; border: 2px solid #00ff99; box-shadow: 0 0 20px #ff6600;"></div>
    <script src="https://3dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
    <script>
      let viewer = $3Dmol.createViewer('viewer', {{ backgroundColor: 'black' }});
      let pdbData = `{pdb_data}`;
      viewer.addModel(pdbData, "pdb");
      viewer.setStyle({{}}, {{cartoon: {{color: 'spectrum'}}}});
      viewer.zoomTo();
      viewer.render();
    </script>
    """

def get_molecule_view(pdb_or_gene: str) -> str:
    """Fetch molecule structure and return 3D viewer HTML."""
    # Try PDB ID first
    pdb_data = fetch_pdb(pdb_or_gene)
    if not pdb_data:
        # Try AlphaFold with UniProt ID or gene name
        pdb_data = fetch_alphafold(pdb_or_gene)
    if pdb_data:
        return render_3dmol(pdb_data)
    return f"<p style='color:red;'>No structure found for {pdb_or_gene}</p>"