import os
import re
import ast
import sys
import yaml
import json
import logging
import pandas as pd
from typing import Dict
from pathlib import Path

def load_system_prompt(system_prompt_file):
    if system_prompt_file and os.path.exists(system_prompt_file):
        with open(system_prompt_file, "r", encoding="utf-8") as f:
            text = f.read()
        match = re.search(r"```\n?(.*?)```", text, re.DOTALL)
        if match:
            return match.group(1).strip()
        return text.strip()
    else:
        raise FileNotFoundError(f"System prompt file not found: {system_prompt_file}")

def load_config(config_path: str = None) -> dict:
    if config_path is None:
        config_path = Path(__file__).resolve().parent.parent / "configs/config.yaml"

    if not Path(config_path).exists():
        raise FileNotFoundError(
            f"Config file not found: {config_path}\n"
        )

    with open(config_path, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f) or {}

    if not isinstance(cfg, dict):
        raise ValueError(
            f"Config file must contain a YAML mapping (dict), got {type(cfg).__name__}. "
        )

    # Normalize None-valued sections to empty dicts so chained .get() never
    # fails with AttributeError (e.g. `llm:` with no sub-keys → None).
    # Recurse into nested dicts so `paths:\n  vector_db:` also gets normalized.
    def _normalize_nulls(d):
        for key in list(d.keys()):
            if d[key] is None:
                d[key] = {}
            elif isinstance(d[key], dict):
                _normalize_nulls(d[key])
    _normalize_nulls(cfg)

    return cfg

def parse_llm_json(raw_text):
    if not raw_text or not raw_text.strip():
        return None, "Empty response"
    text = raw_text.strip()
    # Strip markdown fences
    fence_match = re.search(r"```(?:json)?\s*\n?(.*?)```", text, re.DOTALL)
    if fence_match:
        text = fence_match.group(1).strip()
 
    # Attempt 1: Standard JSON parsing (double quotes)
    try:
        return json.loads(text), None
    except json.JSONDecodeError:
        pass
 
    # Attempt 2: Extract JSON object from surrounding text
    brace_match = re.search(r"\{.*\}", text, re.DOTALL)
    if brace_match:
        extracted = brace_match.group(0)
        try:
            return json.loads(extracted), None
        except json.JSONDecodeError:
            pass
 
        # Attempt 3: Python dict syntax (single quotes) — common with Llama models.
        # ast.literal_eval safely parses Python literals including single-quoted dicts.
        # This handles: {'commodity': 'Coffee', 'intent': 'Crop Nutrition'}
        # and also apostrophes: {'reasoning': "The farmer's crop..."}
        try:
            result = ast.literal_eval(extracted)
            if isinstance(result, dict):
                return result, None
        except (ValueError, SyntaxError):
            pass
 
    return None, f"Could not parse JSON from: {text[:200]}"

def setup_logging(output_dir, log_file):
    log_path = os.path.join(output_dir, log_file)
    # 1. Create a custom module-level logger
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)  # Capture everything at the logger level

    # 2. Create a handler to write logs to a file
    file_handler = logging.FileHandler(log_path, mode="a")
    file_handler.setLevel(logging.INFO)  # File only gets errors and criticals

    # 3. Create a handler to output logs to the console
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)  # Console gets Info, Warning, Error, Critical

    # 4. Define a reusable string format for your log outputs
    formatter = logging.Formatter('%(asctime)s | %(name)s | %(levelname)s | %(message)s')
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)

    # 5. Attach the configured handlers to your custom logger
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)
    return logger

def validate_classification(parsed, classification_fields:Dict):
    errors = []
    if parsed is None:
        return False, ["Parsed result is None"]
    
    for field in classification_fields.keys():
        if field not in parsed:
            errors.append(f"Missing field: {field}")
            continue
        if parsed.get(field) not in classification_fields[field]:
            errors.append(f"Invalid {field}: {parsed.get(field)}")
    return len(errors) == 0, errors

def load_checkpoint(cfg, logger):
    ckpt_path = os.path.join(cfg.output_dir, cfg.checkpoint_file)
    if os.path.exists(ckpt_path):
        df = pd.read_parquet(ckpt_path)
        logger.info(f"RESUMING from checkpoint: {len(df)} queries already classified")
        return df, len(df)
    return None, 0


def save_checkpoint(df, cfg, logger):
    ckpt_path = os.path.join(cfg.output_dir, cfg.checkpoint_file)
    df.to_parquet(ckpt_path, index=False)
    logger.info(f"Checkpoint saved: {len(df)} queries")

def load_queries(cfg, logger):
    # --- Load input data ---
    logger.info(f"Loading data from {cfg.input_path}")
    df_input = pd.read_csv(cfg.input_path) if cfg.input_path.split('.')[-1] == 'csv' else pd.read_parquet(cfg.input_path)

    size = cfg.test_size if cfg.test_size>=0 else df_input.shape[0]
    if cfg.shuffle:
        df_input = df_input.sample(n=size, random_state=42).reset_index(drop=True)

    queries = df_input[cfg.query_column].tolist()
    
    if cfg.classification_type==2:
        assert cfg.response_column is not None, 'Adoption phase analysis requires LLM responses. Specify response column.'
        responses = df_input[cfg.response_column].tolist()
        queries = ['\n\n'.join([f"## Query\n\n{query}", f"### Response\n\n{response}"]) for query, response in zip(queries, responses)]

    return df_input, queries