# swissmodel_tool.py
"""
SWISS-MODEL Repository REST API tool for ToolUniverse.

SWISS-MODEL Repository is a database of annotated 3D protein structure
models generated by the SWISS-MODEL homology-modelling pipeline.
It provides pre-computed models for UniProt entries.

API: https://swissmodel.expasy.org/repository
No authentication required. Free for academic/research use.
"""

import requests
from typing import Dict, Any
from .base_tool import BaseTool
from .tool_registry import register_tool

SWISSMODEL_BASE_URL = "https://swissmodel.expasy.org/repository"


@register_tool("SwissModelTool")
class SwissModelTool(BaseTool):
    """
    Tool for querying the SWISS-MODEL Repository.

    Provides pre-computed protein homology models for UniProt entries,
    including template structures, sequence coverage, quality metrics,
    and download coordinates.

    No authentication required.
    """

    def __init__(self, tool_config: Dict[str, Any]):
        super().__init__(tool_config)
        self.timeout = tool_config.get("timeout", 30)
        self.endpoint_type = tool_config.get("fields", {}).get(
            "endpoint_type", "get_models"
        )

    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Execute the SWISS-MODEL API call."""
        try:
            return self._dispatch(arguments)
        except requests.exceptions.Timeout:
            return {
                "error": f"SWISS-MODEL API request timed out after {self.timeout} seconds"
            }
        except requests.exceptions.ConnectionError:
            return {
                "error": "Failed to connect to SWISS-MODEL API. Check network connectivity."
            }
        except requests.exceptions.HTTPError as e:
            return {"error": f"SWISS-MODEL API HTTP error: {e.response.status_code}"}
        except Exception as e:
            return {"error": f"Unexpected error querying SWISS-MODEL: {str(e)}"}

    def _dispatch(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Route to appropriate endpoint based on config."""
        if self.endpoint_type == "get_models":
            return self._get_models(arguments)
        elif self.endpoint_type == "get_model_summary":
            return self._get_model_summary(arguments)
        else:
            return {"error": f"Unknown endpoint_type: {self.endpoint_type}"}

    def _get_models(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Get all available homology models for a UniProt accession."""
        uniprot_id = arguments.get("uniprot_id", "")
        if not uniprot_id:
            return {
                "error": "uniprot_id parameter is required (e.g., 'P04637' for human p53)"
            }

        url = f"{SWISSMODEL_BASE_URL}/uniprot/{uniprot_id}.json"
        response = requests.get(url, timeout=self.timeout)
        response.raise_for_status()
        raw = response.json()

        result_data = raw.get("result", {})

        # Extract models
        models = []
        for s in result_data.get("structures", []):
            model = {
                "template": s.get("template", ""),
                "method": s.get("method", ""),
                "coverage": s.get("coverage"),
                "from_residue": s.get("from"),
                "to_residue": s.get("to"),
                "created_date": s.get("created_date", ""),
                "provider": s.get("provider", ""),
            }
            # Quality metrics
            if "qmean" in s:
                qmean = s["qmean"]
                if isinstance(qmean, dict):
                    model["qmean_global"] = qmean.get("qmean4_global_score")
                    model["qmean_z_score"] = qmean.get("qmean_z_score")
            # Complex partners if any
            complex_with = s.get("in_complex_with", {})
            if complex_with:
                partners = []
                for chain_id, chain_proteins in complex_with.items():
                    if isinstance(chain_proteins, list):
                        for p in chain_proteins:
                            if isinstance(p, dict):
                                partners.append(
                                    {
                                        "chain": chain_id,
                                        "uniprot_ac": p.get("uniprot_ac", ""),
                                        "description": p.get("description", ""),
                                    }
                                )
                if partners:
                    model["complex_partners"] = partners[:10]
            models.append(model)

        result = {
            "uniprot_id": uniprot_id,
            "sequence_length": result_data.get("sequence_length"),
            "crc64": result_data.get("crc64"),
            "model_count": len(models),
            "models": models[:30],
        }

        return {
            "data": result,
            "metadata": {
                "source": "SWISS-MODEL Repository",
                "api_version": raw.get("api_version", ""),
                "query": uniprot_id,
                "endpoint": "get_models",
            },
        }

    def _get_model_summary(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Get a summary of the best available model for a UniProt accession."""
        uniprot_id = arguments.get("uniprot_id", "")
        if not uniprot_id:
            return {
                "error": "uniprot_id parameter is required (e.g., 'P04637' for human p53)"
            }

        url = f"{SWISSMODEL_BASE_URL}/uniprot/{uniprot_id}.json"
        response = requests.get(url, timeout=self.timeout)
        response.raise_for_status()
        raw = response.json()

        result_data = raw.get("result", {})
        structures = result_data.get("structures", [])

        # Find best model (highest coverage)
        best = None
        best_coverage = 0
        for s in structures:
            cov = s.get("coverage", 0) or 0
            if cov > best_coverage:
                best_coverage = cov
                best = s

        summary = {
            "uniprot_id": uniprot_id,
            "sequence_length": result_data.get("sequence_length"),
            "total_models": len(structures),
        }

        if best:
            summary["best_model"] = {
                "template": best.get("template", ""),
                "method": best.get("method", ""),
                "coverage": best.get("coverage"),
                "from_residue": best.get("from"),
                "to_residue": best.get("to"),
                "created_date": best.get("created_date", ""),
                "coordinates_url": best.get("coordinates", ""),
            }
            if "qmean" in best and isinstance(best["qmean"], dict):
                summary["best_model"]["qmean_global"] = best["qmean"].get(
                    "qmean4_global_score"
                )
        else:
            summary["best_model"] = None

        # Methods summary
        methods = {}
        for s in structures:
            m = s.get("method", "Unknown")
            methods[m] = methods.get(m, 0) + 1
        summary["methods_distribution"] = methods

        return {
            "data": summary,
            "metadata": {
                "source": "SWISS-MODEL Repository",
                "query": uniprot_id,
                "endpoint": "get_model_summary",
            },
        }
