# /services/inference.py

import os
import json
import time
import hashlib
import requests

from typing import Dict, Optional
from threading import Lock

try:
    from dotenv import load_dotenv
    load_dotenv()
except Exception:
    # Environment variables may instead be provided by the service manager.
    pass


# =========================================================
# CONFIGURATION
# =========================================================

BASE_URL = os.getenv(
    "ARDEM_BASE_URL",
    "https://ardemllm.ardem.ai"
)


DEFAULT_TIMEOUT = 60
DEFAULT_POLL_INTERVAL = 5
DEFAULT_MAX_WAIT = 500


# =========================================================
# TOKEN CACHE
# =========================================================

_TOKEN_CACHE = {
    "token": None,
    "created_at": 0
}

_TOKEN_LOCK = Lock()

TOKEN_TTL_SECONDS = 45 * 60


# =========================================================
# UTILITIES
# =========================================================

def compute_md5_of_text(s: str) -> str:
    """
    Normalize text and generate stable MD5.
    """

    if not s:
        return ""

    if not isinstance(s, str):
        s = str(s)

    normalized = " ".join(s.split()).upper()

    return hashlib.md5(
        normalized.encode("utf-8")
    ).hexdigest()


def normalize_response(data: Dict) -> Dict:
    """
    Normalize various possible API response structures.
    """

    if not isinstance(data, dict):
        return {}

    # The document service returns the actual extraction as a JSON string in
    # `extraction_result`. Parse it before considering the response envelope.
    # Returning the envelope used to leak job metadata and non-schema content
    # into downstream state pipelines.
    extraction_result = data.get("extraction_result")
    if isinstance(extraction_result, dict):
        return extraction_result
    if isinstance(extraction_result, str):
        try:
            parsed = json.loads(extraction_result)
            if isinstance(parsed, dict):
                return parsed
        except json.JSONDecodeError:
            pass

    possible_keys = [
        "data",
        "result",
        "results",
        "extracted_data",
        "extraction",
        "payload"
    ]

    for key in possible_keys:
        if key in data and isinstance(data[key], dict):
            return data[key]

    return data


def merge_results(existing: Dict, new: Dict) -> Dict:
    """
    Keep longest non-empty value per field.
    """

    for key, val in new.items():

        if val is None:
            continue

        val = str(val).strip()

        if not val:
            continue

        if (
            key not in existing
            or len(val) > len(str(existing[key]))
        ):
            existing[key] = val

    return existing


# =========================================================
# AUTHENTICATION
# =========================================================

def get_access_token(force_refresh: bool = False) -> str:
    """
    Fetch and cache JWT access token.
    """

    global _TOKEN_CACHE

    # Read at call time: main_db imports this client before its own load_dotenv,
    # and service environments may populate variables after module import.
    username = os.getenv("ARDEM_USERNAME")
    password = os.getenv("ARDEM_PASSWORD")
    if not username or not password:
        raise RuntimeError(
            "ARDEM_USERNAME and ARDEM_PASSWORD must be configured in the environment."
        )

    with _TOKEN_LOCK:

        current_time = time.time()

        existing_token = _TOKEN_CACHE.get("token")
        created_at = _TOKEN_CACHE.get("created_at", 0)

        if (
            existing_token
            and not force_refresh
            and (current_time - created_at) < TOKEN_TTL_SECONDS
        ):
            return existing_token

        url = f"{BASE_URL}/token"

        payload = {
            "username": username,
            "password": password
        }

        headers = {
            "Content-Type": "application/x-www-form-urlencoded"
        }

        response = requests.post(
            url,
            data=payload,
            headers=headers,
            timeout=30
        )

        response.raise_for_status()

        data = response.json()

        token = data.get("access_token")

        if not token:
            raise RuntimeError(
                f"Access token missing: {data}"
            )

        _TOKEN_CACHE = {
            "token": token,
            "created_at": current_time
        }

        return token


# =========================================================
# JOB SUBMISSION
# =========================================================

def submit_document_job(
    token: str,
    s3_input_path: str,
    s3_output_path: str,
    prompt: str,
    project_name: str,
    output_formats=None
) -> str:

    if output_formats is None:
        output_formats = ["json"]  # remove dot

    url = f"{BASE_URL}/process-document"

    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }

    payload = {
        "project_name": project_name,
        "input_file_path": s3_input_path,
        "output_file_path": s3_output_path,
        "prompt": prompt,
        "output_formats": output_formats,
        "notification_webhook": ""
    }

    print("[DEBUG] Submitting payload:")
    print(json.dumps(payload, indent=2))

    response = requests.post(
        url,
        json=payload,
        headers=headers,
        timeout=DEFAULT_TIMEOUT
    )

    if response.status_code >= 400:
        print("[DEBUG] Response content:", response.text)
        response.raise_for_status()

    data = response.json()
    job_id = data.get("job_id")

    if not job_id:
        raise RuntimeError(f"Job submission failed: {data}")

    return job_id

# =========================================================
# POLLING
# =========================================================

def poll_job_result(
    token: str,
    job_id: str,
    poll_interval: int = DEFAULT_POLL_INTERVAL,
    max_wait_time: int = DEFAULT_MAX_WAIT
) -> Dict:

    url = f"{BASE_URL}/result/{job_id}"

    headers = {
        "Authorization": f"Bearer {token}"
    }

    start_time = time.time()

    while True:

        response = requests.get(
            url,
            headers=headers,
            timeout=DEFAULT_TIMEOUT
        )

        if response.status_code == 401:

            token = get_access_token(force_refresh=True)

            headers["Authorization"] = f"Bearer {token}"

            continue

        response.raise_for_status()

        data = response.json()

        status = (
            str(data.get("status", ""))
            .strip()
            .lower()
        )

        print(f"[JOB STATUS] {job_id} → {status}")

        # --------------------------------------------------
        # SUCCESS
        # --------------------------------------------------

        if status == "completed":

            print("\n[FINAL RESPONSE]")
            print(json.dumps(data, indent=2, ensure_ascii=False))

            result = normalize_response(data)

            return result

        # --------------------------------------------------
        # FAILURE
        # --------------------------------------------------

        if status in [
            "failed",
            "failed_ocr",
            "ocr_failed",
            "error"
        ]:

            print("\n[FAILED RESPONSE]")
            print(json.dumps(data, indent=2, ensure_ascii=False))

            raise RuntimeError(
                f"Inference job failed with status '{status}'"
            )

        # --------------------------------------------------
        # TIMEOUT
        # --------------------------------------------------

        elapsed = time.time() - start_time

        if elapsed > max_wait_time:

            print("\n[LAST RESPONSE BEFORE TIMEOUT]")
            print(json.dumps(data, indent=2, ensure_ascii=False))

            raise TimeoutError(
                f"Polling exceeded {max_wait_time}s"
            )

        time.sleep(poll_interval)
        
# =========================================================
# MAIN INFERENCE PIPELINE
# =========================================================

def infer_from_api(
    s3_input_path: str,
    s3_output_path: str,
    prompt: str,
    state: str,
    retries: int = 3
) -> Dict:
    """
    Generic multi-state inference pipeline.
    """

    last_error: Optional[Exception] = None

    for attempt in range(retries):

        try:

            # ---------------------------------------------
            # AUTHENTICATION
            # ---------------------------------------------

            token = get_access_token()

            # ---------------------------------------------
            # JOB SUBMISSION
            # ---------------------------------------------

            job_id = submit_document_job(
                token=token,
                s3_input_path=s3_input_path,
                s3_output_path=s3_output_path,
                prompt=prompt,
                project_name=state
            )

            print(f"[JOB SUBMITTED] {job_id}")

            # ---------------------------------------------
            # RESULT POLLING
            # ---------------------------------------------

            result = poll_job_result(
                token=token,
                job_id=job_id
            )

            if not isinstance(result, dict):

                raise ValueError(
                    f"Invalid response format: {result}"
                )

            # ---------------------------------------------
            # MERGE RESULTS
            # ---------------------------------------------

            combined = {}

            combined = merge_results(
                combined,
                result
            )

            # ---------------------------------------------
            # COMPUTE TXN MD5
            # ---------------------------------------------

            txn_candidates = [
                combined.get("Transaction ID"),
                combined.get("TransactionID"),
                combined.get("Transaction Id"),
                combined.get("Txn"),
                combined.get("transaction_id"),
            ]

            txn_value = next(
                (t for t in txn_candidates if t),
                None
            )

            combined["_txn_md5"] = (
                compute_md5_of_text(txn_value)
                if txn_value
                else ""
            )

            # ---------------------------------------------
            # METADATA
            # ---------------------------------------------

            combined["_job_id"] = job_id
            combined["_state"] = state
            combined["_source"] = "ardem_vllm_pipeline"

            return combined

        except Exception as e:

            last_error = e

            print(
                f"[API ERROR] "
                f"Attempt {attempt+1}/{retries}: {e}"
            )

            time.sleep(3)

    raise RuntimeError(
        f"Pipeline failed after retries: {last_error}"
    )


# =========================================================
# OPTIONAL JSON EXPORT
# =========================================================

def save_json_and_excel(data: Dict, output_path: str):

    os.makedirs(
        os.path.dirname(output_path),
        exist_ok=True
    )

    with open(output_path, "w") as f:

        json.dump(
            data,
            f,
            indent=2,
            ensure_ascii=False
        )


# =========================================================
# TEST RUNNER
# =========================================================

if __name__ == "__main__":

    TEST_S3_INPUT = (
        "s3://bucket/input/sample.pdf"
    )

    TEST_S3_OUTPUT = (
        "s3://bucket/output/sample/"
    )

    TEST_PROMPT = """
    Extract:
    - Transaction ID
    - Process Date
    - Sales Tax

    Return strict JSON only.
    """

    result = infer_from_api(
        s3_input_path=TEST_S3_INPUT,
        s3_output_path=TEST_S3_OUTPUT,
        prompt=TEST_PROMPT,
        state="MO"
    )

    print(json.dumps(result, indent=2))
