import json
import logging
import os
from pathlib import Path
from typing import Any, Dict, Tuple

import pandas as pd
from openai import AzureOpenAI
from dotenv import load_dotenv, find_dotenv

# Load environment variables from .env file
# Search in current directory, parent directory, and examples directory
dotenv_path = find_dotenv(usecwd=True)
if dotenv_path:
    load_dotenv(dotenv_path)
else:
    # Try to load from common locations
    from pathlib import Path
    possible_paths = [
        Path.cwd() / '.env',
        Path.cwd().parent / '.env',
        Path.cwd() / 'examples' / '.env',
        Path(__file__).parent.parent.parent.parent / 'examples' / '.env',
    ]
    for path in possible_paths:
        if path.exists():
            load_dotenv(path)
            break

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def get_openai_client() -> AzureOpenAI:
    """Get configured OpenAI client."""
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        raise ValueError("OPENAI_API_KEY environment variable not set")

    return AzureOpenAI(
        azure_endpoint="https://azurefunctionapi.openai.azure.com/",
        api_key=api_key,
        api_version="2025-01-01-preview",
    )


def generate_all_synthetic_data_batched(
    row_schema: Dict[str, Any], bias_prompt: str, num_rows: int = 50
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    try:
        wrapper_schema = create_wrapper_schema(row_schema, num_rows)
        system_prompt = "You are a data science assistant that generates high-quality, realistic synthetic data based on a provided JSON schema and instructions."
        user_prompt = create_user_prompt(num_rows, bias_prompt, row_schema)

        client = get_openai_client()
        response = client.chat.completions.create(
            model="gpt-4.1-nano",
            messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}],
            tools=[
                {
                    "type": "function",
                    "function": {
                        "name": "generate_datasets",
                        "description": "Generate both neutral and biased datasets.",
                        "parameters": wrapper_schema,
                    },
                }
            ],
            tool_choice={"type": "function", "function": {"name": "generate_datasets"}},
            temperature=0.9,
            max_tokens=8192,
        )

        return process_openai_response(response)

    except Exception as e:
        logger.error(f"OpenAI API call error: {e}")
        return pd.DataFrame(), pd.DataFrame()


def create_wrapper_schema(row_schema: Dict[str, Any], num_rows: int) -> Dict[str, Any]:
    return {
        "type": "object",
        "properties": {
            "neutral_data": {
                "type": "array",
                "items": row_schema,
                "description": f"A diverse, unbiased set of {num_rows} rows.",
            },
            "biased_data": {
                "type": "array",
                "items": row_schema,
                "description": f"A set of {num_rows} rows generated with bias instructions.",
            },
        },
        "required": ["neutral_data", "biased_data"],
    }


def create_user_prompt(num_rows: int, bias_prompt: str, schema: Dict[str, Any]) -> str:
    return f"""Please generate two distinct datasets conforming to the provided schema.
1.  **neutral_data**: Generate {num_rows} diverse and unbiased rows.
2.  **biased_data**: Generate {num_rows} rows, but introduce a strong bias as described here: '{bias_prompt}'.
Schema: ```json
{json.dumps(schema, indent=2)}
```"""


def process_openai_response(response: Any) -> Tuple[pd.DataFrame, pd.DataFrame]:
    try:
        tool_calls = response.choices[0].message.tool_calls
        if not tool_calls:
            logger.error("No tool calls found in OpenAI response")
            return pd.DataFrame(), pd.DataFrame()

        function_args_str = tool_calls[0].function.arguments
        generated_json = json.loads(function_args_str)

        neutral_df = pd.DataFrame(generated_json.get("neutral_data", []))
        biased_df = pd.DataFrame(generated_json.get("biased_data", []))

        return neutral_df, biased_df

    except Exception as e:
        logger.error(f"Error processing OpenAI response: {e}")
        return pd.DataFrame(), pd.DataFrame()
