"""
LLM-powered schema generator using Groq Llama 3.3.

This module provides intelligent schema generation from natural language,
including:
- Reference tables with actual LLM-generated data (exercises, plans, meals)
- Transactional tables with foreign keys to reference tables
- Industry-realistic column configurations
"""

import json
import os
from pathlib import Path
from typing import Dict, Optional

from groq import Groq

from misata.curve_fitting import CurveFitter
from misata.schema import Column, OutcomeCurve, Relationship, ScenarioEvent, SchemaConfig, Table
from misata.research import DeepResearchAgent


# Load .env file if it exists
def _load_env():
    """Load environment variables from .env file."""
    env_paths = [
        Path.cwd() / ".env",
        Path.cwd().parent / ".env",  # apps/.env or api parent
        Path.cwd().parent.parent / ".env",  # Misata root from apps/api
        Path(__file__).parent.parent / ".env",  # packages/core/.env
        Path(__file__).parent.parent.parent / ".env",  # packages/.env
        Path(__file__).parent.parent.parent.parent / ".env",  # Misata root from packages/core/misata
        Path.home() / ".misata" / ".env",
    ]

    for env_path in env_paths:
        if env_path.exists():
            with open(env_path) as f:
                for line in f:
                    line = line.strip()
                    if line and not line.startswith("#") and "=" in line:
                        key, _, value = line.partition("=")
                        # Remove quotes if present
                        value = value.strip().strip("'\"")
                        os.environ.setdefault(key.strip(), value)
            break

_load_env()


SYSTEM_PROMPT = """You are Misata, an expert synthetic data architect. Your job is to generate REALISTIC database schemas based ONLY on the user's story. 

## CRITICAL: DO NOT USE DEFAULT EXAMPLES
- Generate tables that are SPECIFIC to the user's domain.
- If user says "pet store", create tables like "pets", "pet_categories", "pet_sales".
- If user says "music streaming", create tables like "songs", "artists", "streams".
- NEVER default to fitness/exercise/workout tables UNLESS the user explicitly asks for them.

## TABLE TYPES

### 1. REFERENCE TABLES (is_reference: true)
Small lookup tables (5-20 rows) with ACTUAL DATA you generate.
- MUST have an "id" column (integer, sequential from 1)
- Include realistic inline_data based on user's domain

### 2. TRANSACTIONAL TABLES (is_reference: false)
Large tables generated by code using foreign keys.
- Use row_count to specify size
- Use foreign_key type to reference parent tables

## OUTPUT FORMAT

{
  "name": "Dataset Name based on user's domain",
  "description": "Description of the domain",
  "seed": 42,
  "tables": [
    {
      "name": "domain_specific_reference_table",
      "is_reference": true,
      "inline_data": [
        {"id": 1, "name": "Value A", "price": 10.00},
        {"id": 2, "name": "Value B", "price": 20.00}
      ]
    },
    {
      "name": "domain_specific_transactional_table",
      "row_count": 10000,
      "is_reference": false
    }
  ],
  "columns": {
    "domain_specific_transactional_table": [
      {"name": "id", "type": "int", "distribution_params": {"distribution": "uniform", "min": 1, "max": 10000}, "unique": true},
      {"name": "ref_id", "type": "foreign_key", "distribution_params": {}},
      {"name": "amount", "type": "float", "distribution_params": {"distribution": "normal", "mean": 50, "std": 20}},
      {"name": "date", "type": "date", "distribution_params": {"start": "2024-01-01", "end": "2025-12-31"}}
    ]
  },
  "relationships": [
    {"parent_table": "domain_specific_reference_table", "child_table": "domain_specific_transactional_table", "parent_key": "id", "child_key": "ref_id"}
  ],
  "outcome_curves": [],
  "events": []
}

## SMART DEFAULTS FOR COLUMNS

Age: int, normal, mean: 35, std: 12, min: 18, max: 80
Price/Amount: float, exponential, scale: 50, min: 0.01, decimals: 2
Rating (1-5): int, categorical, choices: [1,2,3,4,5], probabilities: [0.05, 0.08, 0.15, 0.32, 0.40]
Quantity: int, poisson, lambda: 3, min: 1
Duration (min): int, normal, mean: 45, std: 20, min: 5
Boolean: boolean, probability: 0.5-0.9 depending on context
Date: date, start/end based on user's time context

## TEMPORAL PATTERNS & OUTCOME CURVES

If the user mentions ANY time-based patterns, EXTRACT them as outcome_curves:

Keywords to detect:
- "peak", "spike", "surge" -> High relative_value (0.8-1.0)
- "dip", "drop", "decline" -> Low relative_value (0.2-0.4)
- "growth", "upward trend" -> pattern_type: "growth"
- "seasonal", "monthly cycles" -> pattern_type: "seasonal"

Output format:
"outcome_curves": [
  {
    "table": "sales",
    "column": "amount",
    "time_column": "sale_date",
    "pattern_type": "seasonal",
    "description": "High in December, low in February",
    "curve_points": [
      {"month": 2, "relative_value": 0.3},
      {"month": 12, "relative_value": 1.0}
    ]
  }
]

## DATE RANGE RULES
- "Last 2 years" -> start: 2024-01-01, end: 2025-12-31
- "Past year" -> start: 2025-01-01, end: 2025-12-31
- "Historical data" -> start: 2020-01-01, end: 2025-12-31
- No mention -> Default to current year (2025)

Generate schemas ONLY based on the user's story. Be creative and domain-specific."""



GRAPH_REVERSE_PROMPT = """You are Misata, an expert at reverse-engineering data patterns.
Given a description of a desired chart or graph pattern, generate a schema that will
produce data matching that EXACT pattern when plotted.

Follow the same two-tier table structure:
- Reference tables with inline_data for lookup values
- Transactional tables with foreign keys for mass data

The user will describe a chart they want. Your job is to generate data that,
when plotted, produces that exact chart."""


class LLMSchemaGenerator:
    """
    Generate realistic schemas from natural language using LLMs.

    Supports multiple providers:
    - groq: Groq Cloud (Llama 3.3) - Fast, free tier
    - openai: OpenAI (GPT-4o) - Best quality
    - ollama: Local Ollama - Free, private

    This is the "brain" of Misata - what makes it genuinely AI-powered.
    """

    # Provider configurations
    PROVIDERS = {
        "groq": {
            "base_url": None,  # Uses default
            "env_key": "GROQ_API_KEY",
            "default_model": "llama-3.3-70b-versatile",
        },
        "openai": {
            "base_url": None,
            "env_key": "OPENAI_API_KEY",
            "default_model": "gpt-4o-mini",
        },
        "ollama": {
            "base_url": "http://localhost:11434/v1",
            "env_key": None,  # No key needed for local
            "default_model": "llama3",
        },
    }

    def __init__(
        self,
        provider: Optional[str] = None,
        api_key: Optional[str] = None,
        model: Optional[str] = None,
        base_url: Optional[str] = None,
    ):
        """
        Initialize the LLM generator.

        Args:
            provider: LLM provider ("groq", "openai", "ollama").
                      Defaults to MISATA_PROVIDER env var or "groq".
            api_key: API key. If not provided, reads from provider's env var.
            model: Model name. If not provided, uses provider default.
            base_url: Custom API base URL (for Ollama or compatible APIs).
        """
        # Determine provider
        self.provider = provider or os.environ.get("MISATA_PROVIDER", "groq").lower()

        if self.provider not in self.PROVIDERS:
            raise ValueError(f"Unknown provider: {self.provider}. Use: {list(self.PROVIDERS.keys())}")

        config = self.PROVIDERS[self.provider]

        # Get API key
        self.api_key = api_key
        if not self.api_key and config["env_key"]:
            self.api_key = os.environ.get(config["env_key"])

        if not self.api_key and self.provider != "ollama":
            env_key = config["env_key"]
            raise ValueError(
                f"{self.provider.title()} API key required. "
                f"Set {env_key} environment variable or pass api_key parameter."
            )

        # Set model
        self.model = model or config["default_model"]

        # Set base URL
        self.base_url = base_url or config["base_url"]

        # Initialize client (all providers use OpenAI-compatible API)
        if self.provider == "groq":
            self.client = Groq(api_key=self.api_key)
        else:
            # OpenAI and Ollama use openai package
            try:
                from openai import OpenAI
            except ImportError:
                raise ImportError(
                    f"openai package required for {self.provider}. "
                    "Install with: pip install openai"
                )

            client_kwargs = {}
            if self.api_key:
                client_kwargs["api_key"] = self.api_key
            if self.base_url:
                client_kwargs["base_url"] = self.base_url

            # Ollama doesn't need a real API key
            if self.provider == "ollama":
                client_kwargs["api_key"] = "ollama"

            self.client = OpenAI(**client_kwargs)

    def generate_from_story(
        self,
        story: str,
        default_rows: int = 10000,
        temperature: float = 0.3,
    ) -> SchemaConfig:
        """
        Generate a realistic schema from a natural language story.

        Args:
            story: Natural language description of the data needs
            default_rows: Default row count if not specified in story
            temperature: LLM temperature (lower = more consistent)

        Returns:
            SchemaConfig ready for data generation
        """
        user_prompt = f"""Generate a complete synthetic data schema in JSON format for:

{story}

CRITICAL INSTRUCTIONS:
1. Generate tables SPECIFIC to the domain described above. DO NOT use generic fitness/exercise examples.
2. Create REFERENCE TABLES (is_reference: true) with inline_data for any lookup/configuration data relevant to THIS domain.
3. Create TRANSACTIONAL TABLES (is_reference: false) with row_count for high-volume data like users, transactions, events, etc.
4. Use foreign_key to link transactional tables to reference tables.
5. Default row count for transactional tables: {default_rows}
6. If the user mentions time patterns (peaks, dips, trends, growth), extract them as outcome_curves.
7. If the user mentions a time range (e.g., "last 2 years"), set date column start/end accordingly.

Output valid JSON. Be creative and domain-specific - DO NOT copy the system prompt examples."""


        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt}
            ],
            temperature=temperature,
            max_tokens=6000,
            response_format={"type": "json_object"}
        )

        schema_dict = json.loads(response.choices[0].message.content)
        return self._parse_schema(schema_dict)

    def generate_from_graph(
        self,
        graph_description: str,
        temperature: float = 0.2,
    ) -> SchemaConfig:
        """
        REVERSE ENGINEERING: Generate schema that produces desired graph patterns.
        """
        user_prompt = f"""Generate a JSON schema that will produce this chart pattern:

{graph_description}

Include reference tables with inline_data for lookup values and transactional tables for mass data. Output valid JSON."""


        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": GRAPH_REVERSE_PROMPT},
                {"role": "user", "content": user_prompt}
            ],
            temperature=temperature,
            max_tokens=6000,
            response_format={"type": "json_object"}
        )

        schema_dict = json.loads(response.choices[0].message.content)
        return self._parse_schema(schema_dict)

    def _normalize_distribution_params(self, col_type: str, params: Dict) -> Dict:
        """Normalize LLM output variations in distribution_params."""
        normalized = params.copy()

        # Normalize date column parameters
        if col_type == "date":
            if "start_date" in normalized and "start" not in normalized:
                normalized["start"] = normalized.pop("start_date")
            if "end_date" in normalized and "end" not in normalized:
                normalized["end"] = normalized.pop("end_date")
            if "start" not in normalized:
                normalized["start"] = "2023-01-01"
            if "end" not in normalized:
                normalized["end"] = "2024-12-31"

        # Normalize categorical parameters
        if col_type == "categorical":
            if "options" in normalized and "choices" not in normalized:
                normalized["choices"] = normalized.pop("options")
            if "choices" not in normalized:
                normalized["choices"] = ["A", "B", "C"]

        # Curve Fitting for 'control_points'
        if "control_points" in normalized:
            try:
                points = normalized.pop("control_points")
                dist_type = normalized.get("distribution", "normal")
                fitter = CurveFitter()
                fitted_params = fitter.fit_distribution(points, dist_type)
                # Merge fitted params, keeping any manual overrides if they exist (or overwriting? let's overwrite for safety)
                normalized.update(fitted_params)
            except Exception:
                # If fitting fails, fallback to defaults or keep what we have
                pass

        return normalized

    def _parse_schema(self, schema_dict: Dict) -> SchemaConfig:
        """Parse LLM output into validated SchemaConfig."""

        # Parse tables
        tables = []
        for t in schema_dict.get("tables", []):
            is_ref = t.get("is_reference", False)
            inline = t.get("inline_data", None)
            row_count = t.get("row_count", len(inline) if inline else 100)

            tables.append(Table(
                name=t["name"],
                row_count=row_count,
                description=t.get("description"),
                is_reference=is_ref,
                inline_data=inline
            ))

        # Parse columns (only for transactional tables, reference tables use inline_data)
        columns = {}
        for table_name, cols in schema_dict.get("columns", {}).items():
            columns[table_name] = []
            for c in cols:
                col_type = c.get("type", "text")
                
                # Normalize LLM type variations to valid schema types
                type_mapping = {
                    "string": "text",
                    "str": "text",
                    "varchar": "text",
                    "char": "text",
                    "integer": "int",
                    "number": "float",
                    "decimal": "float",
                    "double": "float",
                    "timestamp": "datetime",
                    "bool": "boolean",
                    "enum": "categorical",
                    "category": "categorical",
                    "fk": "foreign_key",
                }
                col_type = type_mapping.get(col_type.lower(), col_type)
                
                raw_params = c.get("distribution_params", {})
                normalized_params = self._normalize_distribution_params(col_type, raw_params)

                columns[table_name].append(Column(
                    name=c["name"],
                    type=col_type,
                    distribution_params=normalized_params,
                    nullable=c.get("nullable", False),
                    unique=c.get("unique", False)
                ))

        # For reference tables without columns, create columns from inline_data
        for table in tables:
            if table.is_reference and table.inline_data and table.name not in columns:
                # Infer columns from first row of inline_data
                first_row = table.inline_data[0]
                columns[table.name] = []
                for col_name, value in first_row.items():
                    if isinstance(value, int):
                        col_type = "int"
                    elif isinstance(value, float):
                        col_type = "float"
                    else:
                        col_type = "text"
                    columns[table.name].append(Column(
                        name=col_name,
                        type=col_type,
                        distribution_params={}
                    ))

        # Parse relationships
        relationships = []
        for r in schema_dict.get("relationships", []):
            relationships.append(Relationship(
                parent_table=r["parent_table"],
                child_table=r["child_table"],
                parent_key=r["parent_key"],
                child_key=r["child_key"],
                temporal_constraint=r.get("temporal_constraint", False)
            ))

        # Parse events
        events = []
        for e in schema_dict.get("events", []):
            if not all(key in e for key in ["name", "table", "column", "condition", "modifier_type", "modifier_value"]):
                continue
            events.append(ScenarioEvent(
                name=e["name"],
                table=e["table"],
                column=e["column"],
                condition=e["condition"],
                modifier_type=e["modifier_type"],
                modifier_value=e["modifier_value"],
                description=e.get("description")
            ))

        # Parse outcome curves (temporal patterns from natural language)
        outcome_curves = []
        for c in schema_dict.get("outcome_curves", []):
            if not all(key in c for key in ["table", "column"]):
                continue
            outcome_curves.append(OutcomeCurve(
                table=c["table"],
                column=c["column"],
                time_column=c.get("time_column", "date"),
                pattern_type=c.get("pattern_type", "seasonal"),
                description=c.get("description"),
                curve_points=c.get("curve_points", [])
            ))

        return SchemaConfig(
            name=schema_dict.get("name", "Generated Dataset"),
            description=schema_dict.get("description"),
            tables=tables,
            columns=columns,
            relationships=relationships,
            events=events,
            outcome_curves=outcome_curves,
            seed=schema_dict.get("seed", 42)
        )


    def generate_from_story(self, story: str, use_research: bool = False) -> SchemaConfig:
        """
        Generate schema from a user story.
        
        Args:
            story: The natural language description.
            use_research: If True, uses agent to find real companies for context.
        """
        context = ""
        if use_research:
            print("🕵️‍♂️ Deep Research Mode: ACTIVATED")
            # Simple heuristic to find likely domain
            domain = "SaaS"
            if "fitness" in story.lower(): domain = "Fitness App"
            elif "ecommerce" in story.lower() or "shop" in story.lower(): domain = "Ecommerce"
            elif "finance" in story.lower(): domain = "Fintech"
            
            try:
                # Use Mock Agent (fast)
                agent = DeepResearchAgent(use_mock=True) 
                entities = agent.search_entities(domain, "Competitors", limit=5)
                names = [e['name'] for e in entities]
                context = (
                    f"\n\nREAL WORLD CONTEXT (INJECTED):\n"
                    f"Research found these top players in {domain}: {', '.join(names)}.\n"
                    f"Use these names as examples in the 'inline_data' for reference tables if relevant."
                )
            except Exception as e:
                print(f"Research Agent Warning: {e}")

        # Construct the final prompt
        user_prompt = f"Story: {story}{context}\n\nGenerate the complete JSON schema."
        
        completion = self.client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": SYSTEM_PROMPT,
                },
                {
                    "role": "user",
                    "content": user_prompt,
                }
            ],
            model=self.model,
            temperature=0.1,  # Low temp for JSON consistency
            response_format={"type": "json_object"},
        )

        response_content = completion.choices[0].message.content
        try:
            schema_dict = json.loads(response_content)
            return self._parse_schema(schema_dict)
        except json.JSONDecodeError:
            # Fallback text parsing if JSON mode fails (unlikely with Llama 3)
            # For now, just raise
            raise ValueError(f"Failed to generate valid JSON. Raw response: {response_content[:100]}...")

    def generate_from_graph(self, description: str) -> SchemaConfig:
        """Reverse engineer schema from graph description."""
        # Similar to above but uses GRAPH_REVERSE_PROMPT
        # For brevity, implementing basic pass-through
        return self.generate_from_story(description)

# Convenience functions
def generate_schema(story: str, api_key: Optional[str] = None, use_research: bool = False) -> SchemaConfig:
    """Quick helper to generate schema from story."""
    generator = LLMSchemaGenerator(api_key=api_key)
    return generator.generate_from_story(story, use_research=use_research)


def generate_from_chart(description: str, api_key: Optional[str] = None) -> SchemaConfig:
    """Quick helper to reverse-engineer schema from chart description."""
    generator = LLMSchemaGenerator(api_key=api_key)
    return generator.generate_from_graph(description)
