# coding=utf-8
import json
from unittest.mock import patch

import pytest

from spark_executor.core import connection_store
from spark_executor.models import Connection
from spark_executor.tools import connections, external_jobs


def _fresh_stores():
    """Reset connection store singletons for a single test."""
    store = connection_store.ConnectionStore()
    connection_store.store = store
    connections.store = store
    external_jobs.conn_store = store


def test_get_external_job_logs_returns_tailed():
    _fresh_stores()
    external_jobs.conn_store.save(
        Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
    )
    long_log = "LOG" * 3000
    with patch(
        "spark_executor.tools.external_jobs.get_application_logs",
        return_value=long_log,
    ) as m:
        out = external_jobs.get_external_job_logs(
            application_id="application_1", connection_name="prod", tail_chars=100
        )
    assert out == long_log[-100:]
    args = m.call_args.args
    assert args[0] == "application_1"
    assert args[1].yarn_rm_url == "http://rm:8088"


def test_get_external_job_logs_returns_full_when_short():
    _fresh_stores()
    external_jobs.conn_store.save(
        Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
    )
    short_log = "short log"
    with patch(
        "spark_executor.tools.external_jobs.get_application_logs",
        return_value=short_log,
    ):
        out = external_jobs.get_external_job_logs(
            application_id="application_1", connection_name="prod", tail_chars=5000
        )
    assert out == short_log


def test_get_external_job_logs_raises_when_connection_missing():
    _fresh_stores()
    with pytest.raises(KeyError, match="Connection not found"):
        external_jobs.get_external_job_logs(
            application_id="application_1", connection_name="missing"
        )


def test_get_external_job_status_returns_state():
    _fresh_stores()
    external_jobs.conn_store.save(
        Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
    )
    raw = json.dumps({"app": {"state": "RUNNING"}})
    with patch(
        "spark_executor.tools.external_jobs.get_application_status",
        return_value=("RUNNING", raw),
    ) as m:
        out = external_jobs.get_external_job_status(
            application_id="application_1", connection_name="prod"
        )
    assert out.application_id == "application_1"
    assert out.state == "RUNNING"
    assert out.raw == raw
    args = m.call_args.args
    assert args[0] == "application_1"
    assert args[1].yarn_rm_url == "http://rm:8088"


def test_get_external_job_status_raises_when_connection_missing():
    _fresh_stores()
    with pytest.raises(KeyError, match="Connection not found"):
        external_jobs.get_external_job_status(
            application_id="application_1", connection_name="missing"
        )


def test_get_external_job_result_parses_app_fields():
    _fresh_stores()
    external_jobs.conn_store.save(
        Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
    )
    raw = json.dumps(
        {
            "app": {
                "state": "FINISHED",
                "finalStatus": "SUCCEEDED",
                "diagnostics": "",
                "trackingUrl": "http://rm:8088/proxy/application_1",
                "startedTime": 100,
                "finishedTime": 200,
            }
        }
    )
    with patch(
        "spark_executor.tools.external_jobs.get_application_status",
        return_value=("FINISHED", raw),
    ) as m:
        out = external_jobs.get_external_job_result(
            application_id="application_1", connection_name="prod"
        )
    assert out.application_id == "application_1"
    assert out.state == "FINISHED"
    assert out.final_status == "SUCCEEDED"
    assert out.diagnostics == ""
    assert out.tracking_url == "http://rm:8088/proxy/application_1"
    assert out.started_time == 100
    assert out.finished_time == 200
    args = m.call_args.args
    assert args[0] == "application_1"
    assert args[1].yarn_rm_url == "http://rm:8088"


def test_get_external_job_result_raises_when_connection_missing():
    _fresh_stores()
    with pytest.raises(KeyError, match="Connection not found"):
        external_jobs.get_external_job_result(
            application_id="application_1", connection_name="missing"
        )