#!/usr/bin/env python3
"""
Example Python script to query SynVar and parse the results.
Demonstrates how to use the SynVar API and extract variant information.

API documentation: https://synvar.sibils.org/info
"""

import requests
import xml.etree.ElementTree as ET
import json
from typing import Dict, List, Optional


BASE_URL = "https://synvar.sibils.org/api"


def query_synvar_xml(gene: str, variant: str, level: str = "protein",
                     iso: bool = False, mode: str = "expand") -> str:
    """
    Query the SynVar API and return XML response.

    Args:
        gene: Gene name or reference (e.g., "BRAF", "7")
        variant: Variant description (e.g., "V600E", "c.1799T>A",
                 "NC_000007.14:140753335:A:T" for SPDI, "7:140753336:A:T" for VCF)
        level: Level of variant (protein, transcript, genome, dbsnp, clingen, any)
        iso: Whether to expand to all isoforms
        mode: Processing mode (expand or normalize)

    Returns:
        Response content as string
    """
    params = {
        "ref": gene,
        "variant": variant,
        "level": level,
        "iso": "true" if iso else "false",
        "mode": mode,
    }

    try:
        response = requests.get(BASE_URL, params=params)
        response.raise_for_status()
        return response.text
    except requests.exceptions.RequestException as e:
        print(f"Error downloading: {e}")
        return None


def query_synvar_json(gene: str, variant: str, level: str = "protein",
                      iso: bool = False, mode: str = "expand") -> Optional[Dict]:
    """
    Query the SynVar API and return JSON response.

    Args:
        gene: Gene name or reference
        variant: Variant description
        level: Level of variant (protein, transcript, genome, dbsnp, clingen, any)
        iso: Whether to expand to all isoforms
        mode: Processing mode (expand or normalize)

    Returns:
        Parsed JSON response as dictionary
    """
    params = {
        "ref": gene,
        "variant": variant,
        "level": level,
        "format": "json",
        "iso": "true" if iso else "false",
        "mode": mode,
    }

    try:
        response = requests.get(BASE_URL, params=params)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"Error downloading: {e}")
        return None
    except json.JSONDecodeError as e:
        print(f"Error parsing JSON: {e}")
        return None


def query_synvar_vrs(gene: str, variant: str, level: str = "protein") -> Optional[Dict]:
    """
    Query the SynVar API in VRS format (GA4GH VRS JSON).
    VRS format implies mode=normalize.

    Args:
        gene: Gene name or reference
        variant: Variant description
        level: Level of variant

    Returns:
        Parsed VRS JSON response as dictionary
    """
    params = {
        "ref": gene,
        "variant": variant,
        "level": level,
        "format": "vrs",
    }

    try:
        response = requests.get(BASE_URL, params=params)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"Error downloading: {e}")
        return None
    except json.JSONDecodeError as e:
        print(f"Error parsing JSON: {e}")
        return None


def parse_xml_response(xml_content: str) -> Dict:
    """
    Parse XML response and extract variant information.

    Args:
        xml_content: XML response from SynVar

    Returns:
        Dictionary with parsed variant information
    """
    try:
        root = ET.fromstring(xml_content)
        results = []

        for variant in root.findall('.//variant'):
            variant_info = {
                'valid': variant.get('valid'),
                'mapped': variant.get('mapped'),
                'gene_synonyms': [],
                'protein_synonyms': [],
                'hgvs': [],
                'rsid': None,
                'caid': None,
                'isoforms': []
            }

            for synonym in variant.findall('.//gene-synonym-list/synonym'):
                variant_info['gene_synonyms'].append(synonym.text.strip())

            for synonym in variant.findall('.//protein-synonym-list/synonym'):
                variant_info['protein_synonyms'].append(synonym.text.strip())

            for hgvs in variant.findall('./hgvs'):
                variant_info['hgvs'].append(hgvs.text.strip())

            rsid_elem = variant.find('./rsid')
            if rsid_elem is not None:
                variant_info['rsid'] = rsid_elem.text.strip()

            caid_elem = variant.find('./caid')
            if caid_elem is not None:
                variant_info['caid'] = caid_elem.text.strip()

            # Genome level
            genome_level = variant.find('.//genome-level')
            if genome_level is not None:
                genome_info = {'hgvs': [], 'syntactic_variations': []}
                for hgvs in genome_level.findall('.//hgvs-list/hgvs'):
                    assembly = hgvs.get('assembly', '')
                    genome_info['hgvs'].append({
                        'assembly': assembly,
                        'notation': hgvs.text.strip()
                    })
                for syn_var in genome_level.findall('.//syntactic-variation'):
                    genome_info['syntactic_variations'].append(syn_var.text.strip())
                variant_info['genome_level'] = genome_info

            # Isoforms
            for isoform in variant.findall('.//isoform-list/isoform'):
                isoform_info = {
                    'canonical': isoform.get('canonical-isoform') == 'true',
                    'transcript_level': None,
                    'protein_level': None
                }

                transcript = isoform.find('.//transcript-level')
                if transcript is not None:
                    transcript_info = {'hgvs': [], 'syntactic_variations': []}
                    for hgvs in transcript.findall('.//hgvs'):
                        transcript_info['hgvs'].append(hgvs.text.strip())
                    for syn_var in transcript.findall('.//syntactic-variation'):
                        transcript_info['syntactic_variations'].append(syn_var.text.strip())
                    isoform_info['transcript_level'] = transcript_info

                protein = isoform.find('.//protein-level')
                if protein is not None:
                    protein_info = {'isoform_name': None, 'hgvs': [], 'syntactic_variations': []}
                    iso_name = protein.find('.//isoform-name')
                    if iso_name is not None:
                        protein_info['isoform_name'] = iso_name.text.strip()
                    for hgvs in protein.findall('.//hgvs'):
                        protein_info['hgvs'].append(hgvs.text.strip())
                    for syn_var in protein.findall('.//syntactic-variation'):
                        protein_info['syntactic_variations'].append(syn_var.text.strip())
                    isoform_info['protein_level'] = protein_info

                variant_info['isoforms'].append(isoform_info)

            results.append(variant_info)

        return {'variants': results}

    except ET.ParseError as e:
        print(f"Variant failed to be normalized by SynVar: {e}")
        return None


def print_variant_info(variant_data: Dict):
    """Print variant information in a readable format."""
    for variant in variant_data.get('variants', []):
        print(f"\n{'='*60}")
        print(f"Variant: Valid={variant['valid']}, Mapped={variant['mapped']}")

        if variant['gene_synonyms']:
            print(f"Gene synonyms: {', '.join(variant['gene_synonyms'])}")

        if variant['rsid']:
            print(f"dbSNP: {variant['rsid']}")

        if variant['caid']:
            print(f"ClinGen Allele Registry: {variant['caid']}")

        if variant['hgvs']:
            print(f"Main HGVS: {', '.join(variant['hgvs'])}")

        if 'genome_level' in variant:
            print("\n--- Genome Level ---")
            for hgvs_info in variant['genome_level']['hgvs']:
                print(f"  {hgvs_info['assembly']}: {hgvs_info['notation']}")

        for idx, isoform in enumerate(variant['isoforms'], 1):
            canonical = " (canonical)" if isoform['canonical'] else ""
            print(f"\n--- Isoform {idx}{canonical} ---")

            if isoform['protein_level']:
                print(f"  Protein level:")
                if isoform['protein_level']['isoform_name']:
                    print(f"    Name: {isoform['protein_level']['isoform_name']}")
                for hgvs in isoform['protein_level']['hgvs']:
                    print(f"    HGVS: {hgvs}")
                syn_vars = isoform['protein_level']['syntactic_variations'][:5]
                if syn_vars:
                    print(f"    Syntactic variations (first 5): {', '.join(syn_vars)}")

            if isoform['transcript_level']:
                print(f"  Transcript level:")
                for hgvs in isoform['transcript_level']['hgvs']:
                    print(f"    HGVS: {hgvs}")
                syn_vars = isoform['transcript_level']['syntactic_variations'][:5]
                if syn_vars:
                    print(f"    Syntactic variations (first 5): {', '.join(syn_vars)}")


def main():
    """Example usage of the SynVar API."""

    # Example 1: Protein variant — XML format with parsing
    print("Example 1: BRAF V600E — XML expand mode")
    print("-" * 60)
    xml_content = query_synvar_xml("BRAF", "V600E", "protein")
    if xml_content:
        variant_data = parse_xml_response(xml_content)
        if variant_data:
            print_variant_info(variant_data)

    # Example 2: Protein variant — JSON format
    print("\n\n" + "=" * 60)
    print("Example 2: JAK2 V617F — JSON expand mode")
    print("-" * 60)
    json_data = query_synvar_json("JAK2", "V617F", "protein")
    if json_data:
        variants = json_data.get('variant-list', {}).get('variant', [])
        if not isinstance(variants, list):
            variants = [variants]
        for variant in variants:
            print(f"Valid: {variant.get('@valid')}, Mapped: {variant.get('@mapped')}")
            rsid = variant.get('rsid')
            if rsid:
                print(f"dbSNP: {rsid}")

    # Example 3: Isoform expansion
    print("\n\n" + "=" * 60)
    print("Example 3: TP53 R248W — isoform expansion")
    print("-" * 60)
    json_data = query_synvar_json("TP53", "R248W", "protein", iso=True)
    if json_data:
        variants = json_data.get('variant-list', {}).get('variant', [])
        if not isinstance(variants, list):
            variants = [variants]
        for variant in variants:
            isoforms = variant.get('isoform-list', {}).get('isoform', [])
            if not isinstance(isoforms, list):
                isoforms = [isoforms]
            print(f"Found {len(isoforms)} isoform(s)")
            for idx, isoform in enumerate(isoforms, 1):
                canonical = " (canonical)" if isoform.get('@canonical-isoform') == 'true' else ""
                protein_level = isoform.get('protein-level', {})
                iso_name = protein_level.get('isoform-name', 'Unknown')
                hgvs_list = protein_level.get('hgvs-list', {}).get('hgvs', [])
                if hgvs_list and not isinstance(hgvs_list, list):
                    hgvs_list = [hgvs_list]
                hgvs = hgvs_list[0] if hgvs_list else "N/A"
                print(f"  Isoform {idx}{canonical}: {iso_name} - {hgvs}")

    # Example 4: Normalize mode — returns identifiers only
    print("\n\n" + "=" * 60)
    print("Example 4: BRAF V600E — normalize mode (JSON)")
    print("-" * 60)
    json_data = query_synvar_json("BRAF", "V600E", "protein", mode="normalize")
    if json_data:
        print(json.dumps(json_data, indent=2)[:500])

    # Example 5: VRS format — GA4GH VRS Allele with digest
    print("\n\n" + "=" * 60)
    print("Example 5: BRAF V600E — VRS format (GA4GH)")
    print("-" * 60)
    vrs_data = query_synvar_vrs("BRAF", "V600E", "protein")
    if vrs_data:
        for v in vrs_data.get("variants", []):
            print(f"Gene: {v.get('gene')}")
            print(f"SPDI: {v.get('spdi')}")
            vrs = v.get("vrs", v.get("vrs_allele", {}))
            vrs_id = vrs.get("id", "N/A")
            print(f"VRS ID: {vrs_id}")

    # Example 6: SPDI input
    print("\n\n" + "=" * 60)
    print("Example 6: SPDI input — NC_000007.14:140753335:A:T")
    print("-" * 60)
    xml_content = query_synvar_xml("", "NC_000007.14:140753335:A:T", "genome", mode="normalize")
    if xml_content:
        root = ET.fromstring(xml_content)
        gene = root.findtext('.//gene', '').strip()
        hgvs_g = root.findtext('.//hgvs-genome-grch38', '').strip()
        hgvs_c = root.findtext('.//hgvs-transcript', '').strip()
        hgvs_p = root.findtext('.//hgvs-protein', '').strip()
        print(f"Gene: {gene}")
        print(f"Genomic: {hgvs_g}")
        print(f"Transcript: {hgvs_c}")
        print(f"Protein: {hgvs_p}")


if __name__ == "__main__":
    main()