"""
PTMKB - Get Started
===================

A small client for the Perceptron-PTMKB token API (v2).

1. Create an account on https://perceptronptmkb.lums.edu.pk/v2/login
2. Copy your API token from your profile page (or the user menu).
3. Paste it into ``AUTHORIZATION_TOKEN`` below, or export it as the environment
   variable ``PTMKB_TOKEN``.
4. Run ``python ptmkb_get_started.py`` for an interactive demo, or import the
   functions in your own code.

Requires Python 3.9+ and ``requests`` (``pip install requests``).

Limits
------
* Every call counts against a daily allowance per user (resets at 00:00 UTC).
  Calls over the allowance get HTTP 429.
* Tokens are valid for 30 days. An expired token gets HTTP 403; its detail
  message carries a fresh token, or you can reset it on your profile page.

Legacy paths
------------
Scripts written for PTMKB v1 keep working: the old endpoints are still served
under ``https://perceptronptmkb.lums.edu.pk/ptmkb/api/`` with the v1 response
shapes (``get-ptm-details``, ``get-protein-details``, ``get-available-ptms``,
``get-positional-frequency-matrix``, ``calculate-propensity``). New code should
use the ``/api/v2/ext`` endpoints below. Note the v1 name of the parameter that
selects the modification in ``get-positional-frequency-matrix`` is ``ptm``; the
v1 tutorial script wrongly sent ``selection``.
"""
from __future__ import annotations

import json
import os
from typing import Any, Dict, List, Optional, Union

import requests

BASE_URL = "https://perceptronptmkb.lums.edu.pk/api/v2/ext"
LEGACY_BASE_URL = "https://perceptronptmkb.lums.edu.pk/ptmkb/api"  # v1 shapes, kept for old scripts
TIMEOUT = 30  # seconds

# Please add the token assigned to your account to proceed.
AUTHORIZATION_TOKEN = os.environ.get("PTMKB_TOKEN", "<your_token_here>")

JSON = Union[Dict[str, Any], List[Any], str, int, float, bool, None]


class PTMKBError(RuntimeError):
    """Raised when the API answers with an error or cannot be reached."""

    def __init__(self, message: str, status_code: Optional[int] = None) -> None:
        super().__init__(message)
        self.status_code = status_code


def request(endpoint: str, params: Optional[Dict[str, Any]] = None, base_url: str = BASE_URL) -> JSON:
    """
    Make a GET request to an API endpoint and return the decoded JSON.

    :param endpoint: Path below ``base_url``, e.g. ``"proteins/P04637"``.
    :param params: Query parameters; entries whose value is ``None`` are dropped.
    :param base_url: Override to call the legacy ``/ptmkb/api`` endpoints.
    :raises PTMKBError: on network errors and on HTTP 4xx/5xx answers.

    Example::

        request("proteins/P04637/sites", {"ptm_type": "Phosphorylation"})
    """
    if AUTHORIZATION_TOKEN.startswith("<"):
        raise PTMKBError("Set AUTHORIZATION_TOKEN (or the PTMKB_TOKEN environment variable) first.")
    clean = {k: v for k, v in (params or {}).items() if v is not None}
    try:
        response = requests.get(
            f"{base_url}/{endpoint}",
            params=clean,
            headers={"Authorization": f"Bearer {AUTHORIZATION_TOKEN}"},
            timeout=TIMEOUT,
        )
    except requests.exceptions.RequestException as exc:
        raise PTMKBError(f"Could not reach {endpoint}: {exc}") from exc

    if response.status_code >= 400:
        try:
            detail = response.json().get("detail", response.text)
        except ValueError:
            detail = response.text
        hint = {
            401: " (missing or invalid token)",
            403: " (token expired? the message may contain a fresh token)",
            429: " (daily allowance used up; it resets at 00:00 UTC)",
        }.get(response.status_code, "")
        raise PTMKBError(f"{endpoint}: HTTP {response.status_code}{hint}: {detail}", response.status_code)
    return response.json()


# ---------------------------------------------------------------- v2 endpoints
def get_available_ptms() -> JSON:
    """GET /ptms - modification types in the database."""
    return request("ptms")


def get_protein_details(accession: str) -> JSON:
    """GET /proteins/{acc} - protein record. Accepts UniProt accessions, e.g. ``P04637``."""
    return request(f"proteins/{accession}")


def get_protein_sites(accession: str, ptm_type: Optional[str] = None) -> JSON:
    """GET /proteins/{acc}/sites - modification sites, optionally limited to one PTM type."""
    return request(f"proteins/{accession}/sites", {"ptm_type": ptm_type})


def get_positional_frequency_matrix(ptm: str, residue: str, table: str = "freq") -> JSON:
    """
    GET /ptm-tables/{ptm}/{aa}

    :param ptm: Case-sensitive PTM name from :func:`get_available_ptms`, e.g. ``"Phosphorylation"``.
    :param residue: One-letter code of the modified residue, e.g. ``"S"``.
    :param table: ``"freq"`` (positional frequencies) or ``"log-e"`` (natural-log values).
    """
    if table not in ("freq", "log-e"):
        raise ValueError("table must be 'freq' or 'log-e'")
    return request(f"ptm-tables/{ptm}/{residue}", {"table": table})


def calculate_propensity(ptm: str, subsequence: str) -> JSON:
    """
    GET /propensity/legacy - Log Sum and Log-Log Product of a window.

    :param subsequence: 13 to 21 residues, odd length; the modified residue is in the middle.
    """
    return request("propensity/legacy", {"ptm": ptm, "subsequence": subsequence})


def get_ptm_details(resid: str) -> JSON:
    """GET /resid/{id} - RESID database entry, e.g. ``"AA0039"``."""
    return request(f"resid/{resid}")


def get_propensity_profile(accession: str, ptm: str) -> JSON:
    """GET /proteins/{acc}/propensity/profile - structure-aware v2 score for every eligible residue."""
    return request(f"proteins/{accession}/propensity/profile", {"ptm": ptm})


# ---------------------------------------------------------------- legacy (v1 shapes)
def legacy_get_available_ptms() -> JSON:
    """v1 shape: GET /ptmkb/api/get-available-ptms."""
    return request("get-available-ptms", base_url=LEGACY_BASE_URL)


def legacy_get_protein_details(upid: Optional[str] = None, upac: Optional[str] = None) -> JSON:
    """v1 shape: give either the UniProt identifier (``upid``, e.g. AF9_HUMAN) or the accession (``upac``)."""
    return request("get-protein-details", {"upid": upid, "upac": upac}, base_url=LEGACY_BASE_URL)


def legacy_get_positional_frequency_matrix(ptm: str, residue: str, table: str = "freq") -> JSON:
    """v1 shape. The PTM parameter is called ``ptm`` (the v1 tutorial script sent ``selection``, which was a bug)."""
    return request(
        "get-positional-frequency-matrix",
        {"ptm": ptm, "residue": residue, "table": table},
        base_url=LEGACY_BASE_URL,
    )


# ---------------------------------------------------------------- interactive demo
def _ask(prompt: str, default: str) -> str:
    value = input(f"{prompt} (default: {default!r}): ").strip()
    return value or default


def _show(title: str, data: JSON) -> None:
    print(f"\n{title}:\n{json.dumps(data, indent=2)}")


MENU = (
    "\nChoose a call:\n"
    "  1  Available PTMs\n"
    "  2  Protein details\n"
    "  3  Protein sites\n"
    "  4  Positional frequency matrix\n"
    "  5  Calculate propensity (Log Sum / Log-Log Product)\n"
    "  6  RESID entry\n"
    "  7  Spatial propensity\n"
    "  q  Quit\n"
    "Input: "
)


def main() -> None:
    while True:
        try:
            choice = input(MENU).strip().lower()
            if choice in ("q", "quit", "exit"):
                break
            if choice == "1":
                _show("Available PTMs", get_available_ptms())
            elif choice == "2":
                acc = _ask("UniProt accession", "P04637")
                _show(f"Protein {acc}", get_protein_details(acc))
            elif choice == "3":
                acc = _ask("UniProt accession", "P04637")
                ptm = input("Limit to one PTM type (blank = all): ").strip() or None
                _show(f"Sites of {acc}", get_protein_sites(acc, ptm))
            elif choice == "4":
                # Names are case sensitive: call get_available_ptms() to see how a PTM is spelled.
                ptm = _ask("PTM name", "Phosphorylation")
                residue = _ask("Modified residue", "S")
                table = _ask("Table, 'freq' or 'log-e'", "freq")
                _show("Positional matrix", get_positional_frequency_matrix(ptm, residue, table))
            elif choice == "5":
                ptm = _ask("PTM name", "Phosphorylation")
                sub = _ask("Subsequence (13-21 residues, odd length)", "WKLLPENNVLSPLPSQAMDDW")
                _show("Propensity", calculate_propensity(ptm, sub))
            elif choice == "6":
                resid = _ask("RESID ID", "AA0039")
                _show(f"RESID {resid}", get_ptm_details(resid))
            elif choice == "7":
                acc = _ask("UniProt accession", "P04637")
                ptm = _ask("PTM name", "Phosphorylation")
                _show("Spatial propensity", get_propensity_profile(acc, ptm))
            else:
                print("Please choose one of the listed options.")
        except PTMKBError as exc:
            print(f"\nError: {exc}")
        except (KeyboardInterrupt, EOFError):
            break
    print("\nTerminating program...")


if __name__ == "__main__":
    main()