from pydantic import BaseModel, Field
from typing import List, Dict, Optional, Any
import re
from pathlib import Path
from collections import defaultdict


class EnvVar(BaseModel):
    key: str
    value: str
    comment: Optional[str] = None


class EnvSection(BaseModel):
    name: str  # e.g. "MAIN_DB"
    type: str = "GENERIC"  # e.g. "POSTGRES", "S3"
    vars: List[EnvVar]
    resource_id: Optional[str] = None


class EnvFile(BaseModel):
    sections: List[EnvSection]
    orphan_vars: List[EnvVar] = Field(default_factory=list)


class EnvParser:
    SECTION_REGEX = re.compile(
        r"^(?:#\s*)?\[\s*([^\]]+?)\s*\]\s*(?:\(([^)]+)\))?"
    )  # Matches [NAME], # [NAME], [NAME](TYPE)

    @staticmethod
    def parse(content: str) -> EnvFile:
        sections = []
        orphans = []
        current_section = None

        for line in content.splitlines():
            line = line.strip()
            if not line:
                continue

            # Section Header detection
            sec_match = EnvParser.SECTION_REGEX.match(line)
            if sec_match:
                if current_section:
                    sections.append(current_section)
                
                raw_name = sec_match.group(1).strip()
                type_hint = (
                    sec_match.group(2).strip() if sec_match.lastindex >= 2 else "GENERIC"
                )

                # Handle [TYPE: NAME] or [NAME: TYPE] format
                KNOWN_TYPES = ["POSTGRES", "MYSQL", "CLICKHOUSE", "REDIS", "S3", "MINIO", "OSMNX", "API", "AIRFLOW", "OPENOBSERVE", "MCP"]
                if ":" in raw_name and type_hint == "GENERIC":
                    parts = [p.strip() for p in raw_name.split(":", 1)]
                    # Heuristic: usually it's TYPE: NAME or NAME: TYPE
                    if parts[0].upper() in KNOWN_TYPES:
                        type_hint = parts[0].upper()
                        name = parts[1]
                    elif parts[1].upper() in KNOWN_TYPES:
                        type_hint = parts[1].upper()
                        name = parts[0]
                    else:
                        name = raw_name
                elif type_hint == "GENERIC" and raw_name.upper() in KNOWN_TYPES:
                    type_hint = raw_name.upper()
                    name = raw_name
                else:
                    name = raw_name

                current_section = EnvSection(
                    name=name, type=type_hint or "GENERIC", vars=[]
                )
                continue

            # Var detection
            if "=" in line and not line.startswith("#"):
                key, val = line.split("=", 1)
                # Parse inline comments
                comment = None
                if "#" in val:
                    val, comment = val.split("#", 1)
                    comment = comment.strip()

                var = EnvVar(key=key.strip(), value=val.strip(), comment=comment)

                if current_section:
                    current_section.vars.append(var)
                else:
                    orphans.append(var)

        if current_section:
            sections.append(current_section)

        env_file = EnvFile(sections=sections, orphan_vars=orphans)

        return EnvParser.heuristic_grouping(env_file)

    @staticmethod
    def heuristic_grouping(env_file: EnvFile) -> EnvFile:
        """
        Groups orphan variables into sections based on common prefixes.
        """
        if not env_file.orphan_vars:
            return env_file

        groups = defaultdict(list)
        others = []

        for var in env_file.orphan_vars:
            if "_" in var.key:
                prefix = var.key.split("_", 1)[0]
                groups[prefix].append(var)
            else:
                others.append(var)

        new_sections = list(env_file.sections)

        for prefix, vars_list in groups.items():
            # Check if section already exists
            existing = next((s for s in new_sections if s.name == prefix), None)
            if existing:
                existing.vars.extend(vars_list)
            else:
                # Type Inference
                inferred_type = "GENERIC"
                upper_prefix = prefix.upper()

                # 1. Name-based inference
                KNOWN_TYPES = [
                    "POSTGRES",
                    "MYSQL",
                    "CLICKHOUSE",
                    "REDIS",
                    "S3",
                    "AIRFLOW",
                    "OSMNX",
                    "SENTINEL",
                    "OPENOBSERVE",
                    "MCP",
                    "DASK",
                ]
                if upper_prefix in KNOWN_TYPES:
                    inferred_type = upper_prefix

                # 2. Key-based inference (if still GENERIC)
                if inferred_type == "GENERIC":
                    keys = []
                    for v in vars_list:
                        k = v.key.upper()
                        if k.startswith(upper_prefix + "_"):
                            keys.append(k[len(upper_prefix) + 1 :])
                        else:
                            keys.append(k)

                    keys_set = set(keys)

                    # Signatures
                    if any(
                        "AWS_ACCESS_KEY" in k or "S3_" in k or "FS_KEY" in k or "FS_PATH" in k
                        for k in [v.key.upper() for v in vars_list]
                    ):
                        inferred_type = "S3"
                    elif "HOST" in keys_set and "PORT" in keys_set:
                        if "DB_INDEX" in keys_set or "REDIS_DB" in keys_set:
                            inferred_type = "REDIS"
                        elif (
                            "DATABASE" in keys_set
                            or "DBNAME" in keys_set
                            or "USER" in keys_set
                        ):
                            inferred_type = "POSTGRES"
                        else:
                            pass
                    elif "DB_URL" in keys_set or "DATABASE_URL" in keys_set:
                        inferred_type = "SQLALCHEMY_DATABASE"
                    elif "API_URL" in keys_set and (
                        "USERNAME" in keys_set or "PASSWORD" in keys_set
                    ):
                        inferred_type = "AIRFLOW"

                new_sections.append(
                    EnvSection(name=prefix, type=inferred_type, vars=vars_list)
                )

        return EnvFile(sections=new_sections, orphan_vars=others)

    @staticmethod
    def render(env_file: EnvFile) -> str:
        lines = []

        # Orphans first
        for var in env_file.orphan_vars:
            line = f"{var.key}={var.value}"
            if var.comment:
                line += f" # {var.comment}"
            lines.append(line)

        if lines:
            lines.append("")

        # Sections
        for section in env_file.sections:
            lines.append(f"# [{section.name}]({section.type})")
            for var in section.vars:
                line = f"{var.key}={var.value}"
                if var.comment:
                    line += f" # {var.comment}"
                lines.append(line)
            lines.append("")

        return "\n".join(lines)
