#!/usr/bin/env python3
"""
Redash API Client - Ejecuta queries y obtiene datos desde Redash.

Uso:
    python redash_client.py list-datasources
    python redash_client.py list-queries [--search TEXTO]
    python redash_client.py run-saved QUERY_ID [--format csv|json] [--params '{"key":"value"}']
    python redash_client.py run-adhoc DATASOURCE_ID "SELECT ..." [--format csv|json]
    python redash_client.py get-result QUERY_ID [--format csv|json]
"""

import argparse
import csv
import io
import json
import os
import sys
import time

try:
    import requests
except ImportError:
    print("Error: La libreria 'requests' no esta instalada.", file=sys.stderr)
    print("Instalar con: pip install requests", file=sys.stderr)
    sys.exit(1)


def load_config():
    """Carga la configuracion desde config.json en la carpeta del skill."""
    config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.json")
    if not os.path.exists(config_path):
        print(f"Error: No se encontro el archivo de configuracion en {config_path}", file=sys.stderr)
        print("Crear config.json con: {\"redash_url\": \"https://...\", \"api_key\": \"...\"}", file=sys.stderr)
        sys.exit(1)

    with open(config_path, "r") as f:
        config = json.load(f)

    url = config.get("redash_url", "").rstrip("/")
    api_key = config.get("api_key", "")

    if not url or not api_key:
        print("Error: config.json debe contener 'redash_url' y 'api_key'.", file=sys.stderr)
        sys.exit(1)

    return url, api_key


REDASH_URL = None
API_KEY = None
HEADERS = None
POLL_INTERVAL = 2  # segundos entre cada chequeo del estado de la query
MAX_POLL_ATTEMPTS = 150  # maximo de intentos (5 minutos con intervalo de 2s)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def api_get(endpoint, params=None):
    """Realiza un GET a la API de Redash."""
    url = f"{REDASH_URL}/api/{endpoint}"
    try:
        resp = requests.get(url, headers=HEADERS, params=params, timeout=120)
        resp.raise_for_status()
        return resp.json()
    except requests.exceptions.ConnectionError:
        print(f"Error: No se pudo conectar a {REDASH_URL}. Verificar la URL.", file=sys.stderr)
        sys.exit(1)
    except requests.exceptions.HTTPError as e:
        if resp.status_code == 401:
            print("Error: API Key invalida o sin permisos.", file=sys.stderr)
        elif resp.status_code == 404:
            print(f"Error: Recurso no encontrado en {endpoint}.", file=sys.stderr)
        else:
            print(f"Error HTTP {resp.status_code}: {e}", file=sys.stderr)
        sys.exit(1)


def api_post(endpoint, data=None):
    """Realiza un POST a la API de Redash."""
    url = f"{REDASH_URL}/api/{endpoint}"
    try:
        resp = requests.post(url, headers=HEADERS, json=data, timeout=60)
        resp.raise_for_status()
        return resp.json()
    except requests.exceptions.ConnectionError:
        print(f"Error: No se pudo conectar a {REDASH_URL}. Verificar la URL.", file=sys.stderr)
        sys.exit(1)
    except requests.exceptions.HTTPError as e:
        if resp.status_code == 401:
            print("Error: API Key invalida o sin permisos.", file=sys.stderr)
        elif resp.status_code == 404:
            print(f"Error: Recurso no encontrado en {endpoint}.", file=sys.stderr)
        else:
            print(f"Error HTTP {resp.status_code}: {e}", file=sys.stderr)
            try:
                print(f"Detalle: {resp.text}", file=sys.stderr)
            except Exception:
                pass
        sys.exit(1)


def poll_job(job_id):
    """Espera a que un job de Redash termine y devuelve el query_result_id."""
    for attempt in range(MAX_POLL_ATTEMPTS):
        result = api_get(f"jobs/{job_id}")
        job = result.get("job", result)
        status = job.get("status")

        if status == 3:  # finished
            return job.get("query_result_id")
        elif status == 4:  # error
            error_msg = job.get("error", "Error desconocido")
            print(f"Error al ejecutar la query: {error_msg}", file=sys.stderr)
            sys.exit(1)
        elif status in (1, 2):  # pending / started
            time.sleep(POLL_INTERVAL)
        else:
            print(f"Estado desconocido del job: {status}", file=sys.stderr)
            sys.exit(1)

    print("Error: La query tardo demasiado en ejecutarse (timeout).", file=sys.stderr)
    sys.exit(1)


def format_output(query_result, output_format):
    """Formatea el resultado de una query como CSV o JSON."""
    data = query_result.get("query_result", query_result)
    result_data = data.get("data", data)

    columns = result_data.get("columns", [])
    rows = result_data.get("rows", [])

    if output_format == "json":
        output = {
            "columns": [col.get("name", col.get("friendly_name", "")) for col in columns],
            "rows": rows,
            "row_count": len(rows),
        }
        print(json.dumps(output, indent=2, ensure_ascii=False, default=str))
    else:  # csv
        if not columns or not rows:
            print("Sin resultados.", file=sys.stderr)
            return
        col_names = [col.get("name", col.get("friendly_name", "")) for col in columns]
        writer = csv.DictWriter(sys.stdout, fieldnames=col_names, extrasaction="ignore")
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


# ---------------------------------------------------------------------------
# Comandos
# ---------------------------------------------------------------------------

def cmd_list_datasources():
    """Lista todos los data sources disponibles."""
    datasources = api_get("data_sources")
    if not datasources:
        print("No se encontraron data sources.")
        return

    print(f"{'ID':<6} {'Tipo':<25} {'Nombre'}")
    print("-" * 70)
    for ds in sorted(datasources, key=lambda x: x.get("id", 0)):
        ds_id = ds.get("id", "?")
        ds_type = ds.get("type", "?")
        ds_name = ds.get("name", "?")
        print(f"{ds_id:<6} {ds_type:<25} {ds_name}")


def cmd_list_queries(search=None):
    """Lista queries guardadas, opcionalmente filtradas por texto."""
    params = {"page_size": 50}
    if search:
        params["q"] = search

    result = api_get("queries", params=params)
    queries = result.get("results", result) if isinstance(result, dict) else result

    if not queries:
        print("No se encontraron queries.")
        return

    print(f"{'ID':<8} {'DS ID':<7} {'Nombre'}")
    print("-" * 80)
    for q in queries:
        q_id = q.get("id", "?")
        ds_id = q.get("data_source_id", "?")
        name = q.get("name", "Sin nombre")
        print(f"{q_id:<8} {ds_id:<7} {name}")


def cmd_run_saved(query_id, output_format="csv", params=None):
    """Ejecuta una query guardada y muestra los resultados."""
    data = {}
    if params:
        try:
            parameters = json.loads(params)
            data["parameters"] = parameters
        except json.JSONDecodeError:
            print("Error: Los parametros deben ser un JSON valido.", file=sys.stderr)
            sys.exit(1)

    # Primero intentar obtener resultado existente
    result = api_post(f"queries/{query_id}/results", data=data if data else None)

    if "job" in result:
        # La query se puso en cola, esperar
        job_id = result["job"]["id"]
        print(f"Query en ejecucion (job {job_id}), esperando...", file=sys.stderr)
        query_result_id = poll_job(job_id)
        result = api_get(f"query_results/{query_result_id}")

    format_output(result, output_format)


def cmd_run_adhoc(datasource_id, query_text, output_format="csv"):
    """Ejecuta una query SQL ad-hoc contra un data source."""
    data = {
        "data_source_id": int(datasource_id),
        "query": query_text,
        "max_age": 0,  # forzar ejecucion fresca
    }

    result = api_post("query_results", data=data)

    if "job" in result:
        job_id = result["job"]["id"]
        print(f"Query en ejecucion (job {job_id}), esperando...", file=sys.stderr)
        query_result_id = poll_job(job_id)
        result = api_get(f"query_results/{query_result_id}")

    format_output(result, output_format)


def cmd_get_result(query_id, output_format="csv"):
    """Obtiene el ultimo resultado de una query guardada (sin re-ejecutar)."""
    result = api_get(f"queries/{query_id}/results")
    format_output(result, output_format)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Cliente de Redash - Ejecutar queries y obtener datos",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    subparsers = parser.add_subparsers(dest="command", help="Comando a ejecutar")

    # list-datasources
    subparsers.add_parser("list-datasources", help="Listar data sources disponibles")

    # list-queries
    lq = subparsers.add_parser("list-queries", help="Listar queries guardadas")
    lq.add_argument("--search", "-s", help="Filtrar queries por texto")

    # run-saved
    rs = subparsers.add_parser("run-saved", help="Ejecutar una query guardada")
    rs.add_argument("query_id", type=int, help="ID de la query en Redash")
    rs.add_argument("--format", "-f", choices=["csv", "json"], default="csv", help="Formato de salida")
    rs.add_argument("--params", "-p", help='Parametros como JSON: \'{"key": "value"}\'')

    # run-adhoc
    ra = subparsers.add_parser("run-adhoc", help="Ejecutar SQL ad-hoc")
    ra.add_argument("datasource_id", type=int, help="ID del data source")
    ra.add_argument("query", help="Query SQL a ejecutar")
    ra.add_argument("--format", "-f", choices=["csv", "json"], default="csv", help="Formato de salida")

    # get-result
    gr = subparsers.add_parser("get-result", help="Obtener ultimo resultado de una query")
    gr.add_argument("query_id", type=int, help="ID de la query en Redash")
    gr.add_argument("--format", "-f", choices=["csv", "json"], default="csv", help="Formato de salida")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(1)

    global REDASH_URL, API_KEY, HEADERS
    REDASH_URL, API_KEY = load_config()
    HEADERS = {"Authorization": f"Key {API_KEY}", "Content-Type": "application/json"}

    if args.command == "list-datasources":
        cmd_list_datasources()
    elif args.command == "list-queries":
        cmd_list_queries(search=args.search)
    elif args.command == "run-saved":
        cmd_run_saved(args.query_id, output_format=args.format, params=args.params)
    elif args.command == "run-adhoc":
        cmd_run_adhoc(args.datasource_id, args.query, output_format=args.format)
    elif args.command == "get-result":
        cmd_get_result(args.query_id, output_format=args.format)


if __name__ == "__main__":
    main()
