import json
import logging
import litellm
from datetime import datetime
from importlib.resources import files
from pathlib import Path
import random

from dotenv import load_dotenv

from opensyndrome.artifacts import get_schema_filepath
from opensyndrome.schema import OpenSyndromeCaseDefinitionSchema
from opensyndrome.providers import (
    DEFAULT_MODEL,
    DEFAULT_PROVIDER,
    build_model_string,
    get_instructor_client,
    get_litellm_kwargs,
)

load_dotenv()
logger = logging.getLogger(__name__)


def load_examples(examples_dir: Path, random_k=None):
    json_definitions = {}
    for raw_json in examples_dir.glob("**/*"):
        if not raw_json.name.endswith(".json"):
            continue
        if raw_json.read_text(encoding="utf-8") != "":
            content = json.loads(raw_json.read_text(encoding="utf-8"))
            if content:
                json_definitions[raw_json.stem] = content

    definitions = list(json_definitions.values())
    if random_k:
        definitions = random.sample(definitions, random_k)
    examples = "\n".join(
        f"- {json.dumps(_definition)}" for _definition in definitions if _definition
    )
    return examples


PROMPT_TO_MACHINE_READABLE_FORMAT = """
You are an expert in creating standardized case definition JSONs for medical syndromes.
Generate a JSON that strictly follows this JSON schema, using the provided example documents as reference.

Strict Rules:
- ONLY use symptoms explicitly mentioned in the input text
- ONLY use criteria explicitly mentioned in the schema
- Generate JSON matching the provided schema exactly
- Do not add any information not in the source text
- Use logical operators to capture text's precise meaning
- If text is ambiguous, minimize assumptions

Example documents to reference:
{examples}

Input: {human_readable_definition}

Expected Output Format:
- Use {language} language
- JSON matching provided schema only
- Criteria reflecting ONLY input text
- No additional professional judgment or external information
"""
PROMPT_TO_HUMAN_READABLE_FORMAT = """
You are an public health expert in creating standardized case definition.
Generate a human-readable definition from the provided JSON using only
the information provided there.

Expected output format:
- Text in narrative form
- Clear, concise, and easy to understand text
- No additional professional judgment or external information
- Write in {language}

{machine_readable_definition}
"""


def _add_first_level_required_fields(schema: dict, definition: dict):
    """Add mandatory fields and empty values as placeholders."""
    default_values = {
        "string": "",
        "array": [],
        "object": {},
        "integer": 0,
    }
    missing_fields = set(schema["required"]) - set(definition.keys())
    for field in missing_fields:
        definition[field] = default_values.get(schema["properties"][field]["type"])
    return definition


def _fill_automatic_fields(
    machine_readable_definition: dict, human_readable_definition: str
):
    machine_readable_definition["human_readable_definition"] = human_readable_definition
    machine_readable_definition["published_in"] = (
        "https://opensyndrome.org/definitions/<replace-url>"
    )
    machine_readable_definition["published_at"] = str(
        datetime.now().strftime("%Y-%m-%dT%H:%M:%SZ")
    )
    machine_readable_definition["published_by"] = []
    machine_readable_definition["status"] = "draft"
    machine_readable_definition["open_syndrome_version"] = (
        "1.0.0"  # TODO get this version from definition repo
    )
    machine_readable_definition["references"] = [
        {"citation": "", "url": ""}
    ]  # to be filled by the user
    schema = json.loads(get_schema_filepath().read_text())
    machine_readable_definition = _add_first_level_required_fields(
        schema, machine_readable_definition
    )
    return machine_readable_definition


def generate_machine_readable_format(
    human_readable_definition,
    model=DEFAULT_MODEL,
    language="American English",
    provider=DEFAULT_PROVIDER,
):
    if not human_readable_definition:
        raise ValueError("Human-readable definition cannot be empty.")

    examples_dir = files("opensyndrome").joinpath("examples")
    examples = load_examples(examples_dir, 3)
    formatted_prompt = PROMPT_TO_MACHINE_READABLE_FORMAT.format(
        examples=examples,
        human_readable_definition=human_readable_definition,
        language=language,
    )

    client = get_instructor_client(provider)
    model_string = build_model_string(provider, model)
    instance = client.chat.completions.create(
        model=model_string,
        messages=[{"role": "user", "content": formatted_prompt}],
        response_model=OpenSyndromeCaseDefinitionSchema,
        temperature=0,
        **get_litellm_kwargs(provider),
    )

    return _fill_automatic_fields(
        instance.model_dump(exclude_none=True, by_alias=True, mode="json"),
        human_readable_definition,
    )


def _exclude_metadata_fields(definition: dict):
    """Exclude metadata fields from the definition."""
    definition_fields = [
        "inclusion_criteria",
        "exclusion_criteria",
        "target_public_health_threats",
    ]
    return {key: value for key, value in definition.items() if key in definition_fields}


def generate_human_readable_format(
    machine_readable_definition,
    model=DEFAULT_MODEL,
    language="American English",
    provider=DEFAULT_PROVIDER,
):
    if not machine_readable_definition:
        raise ValueError("Machine-readable definition cannot be empty.")

    formatted_prompt = PROMPT_TO_HUMAN_READABLE_FORMAT.format(
        language=language,
        machine_readable_definition=_exclude_metadata_fields(
            machine_readable_definition
        ),
    )
    model_string = build_model_string(provider, model)
    response = litellm.completion(
        model=model_string,
        messages=[{"role": "user", "content": formatted_prompt}],
        temperature=0,
        **get_litellm_kwargs(provider),
    )
    return response.choices[0].message.content
