import pandas as pd
from typing import Union

from src.generator import LLMGenerator, ModelConfig
from src.logger import get_logger

logger = get_logger(__name__)


def parse_judgement(raw: str) -> Union[dict, None]:
    """
    Expects a judgement string as an output from the judgement generator.
    Parses the judgement string to JSON format.
    
    Expected input format:
    {"accuracy": int, "explanation": str}
    {"fluency": int, "explanation": str}
    {"grammar": int, "explanation": str}
    {"style": int, "explanation": str}
    {"completeness": int, "explanation": str}
    """
    def parse_line(line: str) -> tuple[str, int, str]:
        content = line.strip("{}")
        parts = content.split(",", 1)
        if len(parts) != 2:
            raise ValueError(f"Invalid judgement format: {line}")
        
        key, score = parts[0].strip().split(":")
        key = key.strip().strip('"')
        score = int(score.strip())

        explanation = parts[1].strip().split(":", 1)[1].strip().strip('"')

        return key, score, explanation

    result = {}
    
    try:
        for line in raw.splitlines():
            line = line.strip()
            if not line or not line.startswith("{") or not line.endswith("}"):
                continue

            key, score, explanation = parse_line(line)
            result[key] = {
                "score": score,
                "explanation": explanation
            }

        if len(result.keys()) != 5:
            logger.warning("Judgement does not contain all required fields: %s", raw)
    except:
        logger.error("Error parsing judgement: %s", raw)

    return result if result else None


def make_judgements(data: pd.DataFrame, config: ModelConfig) -> pd.DataFrame:
    logger.info("Setting up the judgement generator with the provided configuration.")
    logger.info(f"Using the following configuration: {repr(config)}")

    judge = LLMGenerator(config)

    logger.info("Starting generation of judgements.")

    prompt = lambda sample: config.prompt_template.format(en=sample['en'], pl=sample['pl']).strip()
    
    items = data.reset_index().explode('conversations')
    items['pl'] = items['conversations'].apply(lambda x: x['pl'])
    items['en'] = items['conversations'].apply(lambda x: x['en'])

    samples = [{'en': item['en'], 'pl': item['pl']} for _, item in items.iterrows()]
    items['judgement'] = judge.generate(samples, prompt)
    
    conversations = [[] for _ in range(len(data))]
    for _, row in items.iterrows():
        item = row['conversations']
        judgement = parse_judgement(row['judgement'])
        item['judgement'] = judgement
        conversations[row['index']].append(item)

    judgements_data = data.copy()
    judgements_data['conversations'] = conversations

    logger.info("Judgements generation completed successfully.")
    return judgements_data

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str, help="Path to the input file", required=True)
    parser.add_argument("--config", type=str, help="Path to the configuration file", required=True)
    parser.add_argument("--output", type=str, help="Path to the output file", required=True)

    args = parser.parse_args()

    data = pd.read_json(args.input, lines=args.input.endswith(".jsonl"))
    config = ModelConfig.from_yaml(args.config)

    judgements = make_judgements(data, config)
    use_lines = args.output.endswith(".jsonl")
    judgements.to_json(args.output, orient="records", lines=use_lines, indent=2 if not use_lines else None)

    logger.info("Judgements saved to %s", args.output)