"""
Stage 3: PDB Structure & UniProt ID Fetch
Retrieves PDB structure IDs and UniProt accession numbers for target proteins.
Includes file-based caching to avoid redundant API calls.
"""
import requests
import json
import logging
import os
from datetime import datetime
from typing import List, Dict, Optional

logger = logging.getLogger(__name__)


class PdbStructurePipeline:
    def __init__(self, api_url: str = "https://rest.uniprot.org/uniprotkb/search",
                 cache_dir: str = None):
        self.api_url = api_url
        if cache_dir is None:
            cache_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "data")
        self.cache_dir = cache_dir
        os.makedirs(self.cache_dir, exist_ok=True)

    def _get_cache_path(self, symbol: str) -> str:
        return os.path.join(self.cache_dir, f"cache_pdb_{symbol.lower()}.json")

    def _load_from_cache(self, symbol: str) -> Optional[Dict]:
        cache_path = self._get_cache_path(symbol)
        if os.path.exists(cache_path):
            try:
                with open(cache_path, "r") as f:
                    data = json.load(f)
                if data.get("symbol") == symbol:
                    logger.info(f"  📦 Loaded PDB data from cache for: {symbol}")
                    return data
            except (json.JSONDecodeError, IOError) as e:
                logger.warning(f"  ⚠️ Failed to read cache for {symbol}: {e}")
        return None

    def _save_to_cache(self, symbol: str, data: Dict):
        cache_path = self._get_cache_path(symbol)
        data["fetched_at"] = datetime.utcnow().isoformat()
        try:
            with open(cache_path, "w") as f:
                json.dump(data, f, indent=2)
        except IOError as e:
            logger.warning(f"  ⚠️ Failed to write cache for {symbol}: {e}")

    def fetch_pdb_ids(self, targets: List[Dict]) -> List[Dict]:
        logger.info(f"Fetching PDB structures for {len(targets)} targets...")
        enriched = []
        for t in targets:
            symbol = t["symbol"]
            cached = self._load_from_cache(symbol)
            if cached is not None:
                t["pdb_ids"] = cached.get("pdb_ids", [])
                t["uniprot_id"] = cached.get("uniprot_id", "")
                enriched.append(t)
                continue
            try:
                params = {
                    "query": f"{symbol} AND taxonomy_id:9606",
                    "format": "json",
                    "size": 1,
                }
                response = requests.get(self.api_url, params=params, timeout=30)
                response.raise_for_status()
                data = response.json()
                uniprot_id = ""
                pdb_ids = []
                if data.get("results") and len(data["results"]) > 0:
                    result = data["results"][0]
                    uniprot_id = result.get("primaryAccession", "")
                    xrefs = result.get("uniProtKBCrossReferences", [])
                    for xref in xrefs:
                        if xref.get("database") == "PDB":
                            pdb_id = xref.get("id")
                            if pdb_id:
                                pdb_ids.append(pdb_id)
                    pdb_ids = pdb_ids[:3]
                t["pdb_ids"] = pdb_ids
                t["uniprot_id"] = uniprot_id
                if pdb_ids:
                    logger.info(f"  ✅ Found {len(pdb_ids)} PDB IDs for: {symbol}")
                else:
                    logger.warning(f"  ⚠️ No PDB structures found for: {symbol}")
                self._save_to_cache(symbol, {
                    "symbol": symbol,
                    "uniprot_id": uniprot_id,
                    "pdb_ids": pdb_ids,
                })
            except requests.exceptions.RequestException as e:
                logger.warning(f"  ⚠️ Failed to fetch PDB for {symbol}: {e}")
                t["pdb_ids"] = []
                t["uniprot_id"] = ""
            except Exception as e:
                logger.warning(f"  ⚠️ Failed to fetch PDB for {symbol}: {e}")
                t["pdb_ids"] = []
                t["uniprot_id"] = ""
            enriched.append(t)
        logger.info(f"✅ PDB fetch complete for {len(enriched)}/{len(targets)} targets")
        return enriched