"""AI-powered AISA compliance checker.

Uses an LLM agent to semantically analyze a project's code against
AISA layers, providing deeper analysis than pattern matching alone.
Supports OpenAI, Anthropic, and Google Gemini models.
"""

import json
import os
from pathlib import Path
from dataclasses import dataclass

from ..scanner.analyzer import ScanResult, LayerResult, CriterionResult
from ..layers import ALL_LAYERS, ALL_CONTRACTS


SKILLS_PATH = Path(__file__).parent.parent / "skills.md"

# Max characters of source code to send to the LLM
MAX_SOURCE_CHARS = 80_000
# Max file size to include
MAX_FILE_SIZE = 50_000

SUPPORTED_PROVIDERS = {
    "openai": ["OPENAI_API_KEY"],
    "anthropic": ["ANTHROPIC_API_KEY"],
    "google": ["GOOGLE_API_KEY", "GEMINI_API_KEY"],
}


@dataclass
class AICheckConfig:
    """Configuration for the AI checker."""
    provider: str = "openai"
    model: str = ""
    api_key: str = ""
    temperature: float = 0.1

    def resolve(self) -> "AICheckConfig":
        """Auto-detect provider and API key from environment."""
        if self.api_key:
            return self

        # Try each provider in order
        for provider, env_vars in SUPPORTED_PROVIDERS.items():
            for env_var in env_vars:
                key = os.environ.get(env_var, "")
                if key:
                    self.provider = provider
                    self.api_key = key
                    if not self.model:
                        self.model = _default_model(provider)
                    return self

        return self


def _default_model(provider: str) -> str:
    defaults = {
        "openai": "gpt-4o",
        "anthropic": "claude-sonnet-4-20250514",
        "google": "gemini-2.0-flash",
    }
    return defaults.get(provider, "gpt-4o")


def ai_check_project(project_path: Path, config: AICheckConfig | None = None) -> ScanResult:
    """Run AI-powered AISA compliance check on a project.

    Args:
        project_path: Root directory of the project.
        config: LLM configuration (auto-detects if not provided).

    Returns:
        ScanResult with AI-generated scores and evidence.
    """
    if config is None:
        config = AICheckConfig()
    config = config.resolve()

    if not config.api_key:
        raise RuntimeError(
            "No API key found. Set one of: OPENAI_API_KEY, ANTHROPIC_API_KEY, or GOOGLE_API_KEY\n"
            "Or pass --provider and --api-key flags."
        )

    # 1. Collect project context
    project_context = _build_project_context(project_path)

    # 2. Load skills prompt
    skills_prompt = SKILLS_PATH.read_text()

    # 3. Build the analysis prompt
    analysis_prompt = _build_analysis_prompt(skills_prompt, project_context)

    # 4. Call the LLM
    print(f"  Analyzing with {config.provider}/{config.model}...")
    response_text = _call_llm(analysis_prompt, config)

    # 5. Parse response into ScanResult
    result = _parse_ai_response(response_text, str(project_path))

    return result


def _build_project_context(project_path: Path) -> str:
    """Build a context string from the project's source files."""
    parts: list[str] = []
    total_chars = 0

    skip_dirs = {".git", ".venv", "venv", "env", "node_modules", "__pycache__",
                 ".tox", ".mypy_cache", ".pytest_cache", "dist", "build", ".eggs"}

    # Prioritize key files
    priority_files = [
        "README.md", "README.rst", "pyproject.toml", "setup.py",
        "requirements.txt", "Dockerfile", ".env.example",
    ]

    for pf in priority_files:
        fp = project_path / pf
        if fp.exists() and fp.is_file():
            try:
                content = fp.read_text(errors="ignore")[:MAX_FILE_SIZE]
                header = f"\n--- FILE: {fp.relative_to(project_path)} ---\n"
                parts.append(header + content)
                total_chars += len(content)
            except (PermissionError, OSError):
                continue

    # Then collect Python files
    py_files: list[Path] = []
    try:
        for py_file in sorted(project_path.rglob("*.py")):
            if any(part in skip_dirs or part.startswith(".") for part in py_file.parts):
                continue
            py_files.append(py_file)
    except (PermissionError, OSError):
        pass

    # Sort by likely importance: shorter paths first, then test files last
    def sort_key(p: Path) -> tuple:
        is_test = "test" in p.name.lower()
        return (is_test, len(p.parts), p.name)

    py_files.sort(key=sort_key)

    for py_file in py_files:
        if total_chars >= MAX_SOURCE_CHARS:
            parts.append(f"\n--- TRUNCATED: {len(py_files) - len(parts)} more Python files not shown ---\n")
            break
        try:
            content = py_file.read_text(errors="ignore")
            if len(content) > MAX_FILE_SIZE:
                content = content[:MAX_FILE_SIZE] + "\n# ... truncated ..."
            header = f"\n--- FILE: {py_file.relative_to(project_path)} ---\n"
            parts.append(header + content)
            total_chars += len(content)
        except (PermissionError, OSError):
            continue

    # Also check for CI/CD, config files
    config_patterns = ["*.yaml", "*.yml", "*.toml", "*.json"]
    ci_dirs = [".github/workflows", ".gitlab-ci"]

    for ci_dir in ci_dirs:
        ci_path = project_path / ci_dir
        if ci_path.exists():
            for f in sorted(ci_path.rglob("*")):
                if f.is_file() and total_chars < MAX_SOURCE_CHARS:
                    try:
                        content = f.read_text(errors="ignore")[:MAX_FILE_SIZE]
                        header = f"\n--- FILE: {f.relative_to(project_path)} ---\n"
                        parts.append(header + content)
                        total_chars += len(content)
                    except (PermissionError, OSError):
                        continue

    # Project structure
    structure = _get_project_structure(project_path, skip_dirs)
    parts.insert(0, f"--- PROJECT STRUCTURE ---\n{structure}\n")

    return "\n".join(parts)


def _get_project_structure(project_path: Path, skip_dirs: set) -> str:
    """Get a tree-like view of the project structure."""
    lines: list[str] = []
    count = 0
    max_entries = 200

    for item in sorted(project_path.rglob("*")):
        if count >= max_entries:
            lines.append(f"... and more files (truncated at {max_entries})")
            break
        if any(part in skip_dirs or part.startswith(".") for part in item.relative_to(project_path).parts):
            continue
        rel = item.relative_to(project_path)
        if item.is_file():
            size = item.stat().st_size
            lines.append(f"  {rel} ({size:,} bytes)")
            count += 1

    return "\n".join(lines)


def _build_analysis_prompt(skills: str, context: str) -> str:
    """Build the full prompt for the LLM."""
    return f"""{skills}

---

## Project to Analyze

{context}

---

Now analyze this project against all AISA criteria. Return ONLY valid JSON matching the output format specified above. Score every criterion from L1.1 through C4.3."""


def _call_llm(prompt: str, config: AICheckConfig) -> str:
    """Call the LLM provider and return the response text."""
    if config.provider == "openai":
        return _call_openai(prompt, config)
    elif config.provider == "anthropic":
        return _call_anthropic(prompt, config)
    elif config.provider == "google":
        return _call_google(prompt, config)
    else:
        raise ValueError(f"Unsupported provider: {config.provider}")


def _call_openai(prompt: str, config: AICheckConfig) -> str:
    try:
        from openai import OpenAI
    except ImportError:
        raise ImportError("OpenAI package required: pip install openai")

    client = OpenAI(api_key=config.api_key)
    response = client.chat.completions.create(
        model=config.model,
        messages=[
            {"role": "system", "content": "You are an expert AI systems architect performing an AISA compliance assessment. Return only valid JSON."},
            {"role": "user", "content": prompt},
        ],
        temperature=config.temperature,
        max_tokens=8000,
        response_format={"type": "json_object"},
    )
    return response.choices[0].message.content or ""


def _call_anthropic(prompt: str, config: AICheckConfig) -> str:
    try:
        import anthropic
    except ImportError:
        raise ImportError("Anthropic package required: pip install anthropic")

    client = anthropic.Anthropic(api_key=config.api_key)
    response = client.messages.create(
        model=config.model,
        max_tokens=8000,
        messages=[
            {"role": "user", "content": f"You are an expert AI systems architect performing an AISA compliance assessment. Return only valid JSON.\n\n{prompt}"},
        ],
        temperature=config.temperature,
    )
    return response.content[0].text


def _call_google(prompt: str, config: AICheckConfig) -> str:
    try:
        import google.generativeai as genai
    except ImportError:
        raise ImportError("Google GenAI package required: pip install google-generativeai")

    genai.configure(api_key=config.api_key)
    model = genai.GenerativeModel(config.model)
    response = model.generate_content(
        f"You are an expert AI systems architect performing an AISA compliance assessment. Return only valid JSON.\n\n{prompt}",
        generation_config=genai.types.GenerationConfig(
            temperature=config.temperature,
            max_output_tokens=8000,
        ),
    )
    return response.text


def _parse_ai_response(response_text: str, project_path: str) -> ScanResult:
    """Parse the LLM's JSON response into a ScanResult."""
    # Extract JSON from response (handle markdown code blocks)
    json_text = response_text.strip()
    if json_text.startswith("```"):
        # Remove markdown code block
        lines = json_text.split("\n")
        json_lines = []
        in_block = False
        for line in lines:
            if line.startswith("```") and not in_block:
                in_block = True
                continue
            elif line.startswith("```") and in_block:
                break
            elif in_block:
                json_lines.append(line)
        json_text = "\n".join(json_lines)

    try:
        data = json.loads(json_text)
    except json.JSONDecodeError:
        # Try to find JSON object in the response
        start = json_text.find("{")
        end = json_text.rfind("}") + 1
        if start >= 0 and end > start:
            try:
                data = json.loads(json_text[start:end])
            except json.JSONDecodeError:
                return _empty_result(project_path, "Failed to parse AI response as JSON")
        else:
            return _empty_result(project_path, "No JSON found in AI response")

    scores = data.get("scores", {})
    frameworks = data.get("frameworks_detected", [])

    # Build framework dict
    fw_dict = {fw: 1.0 for fw in frameworks} if isinstance(frameworks, list) else {}

    # Build layer results
    layer_results = []
    for layer in ALL_LAYERS:
        criteria_results = []
        for criterion in layer.criteria:
            score_data = scores.get(criterion.id, {})
            if isinstance(score_data, dict):
                score_val = score_data.get("score", "no")
                evidence_text = score_data.get("evidence", "")
                file_ref = score_data.get("file", "")
            elif isinstance(score_data, str):
                score_val = score_data
                evidence_text = ""
                file_ref = ""
            else:
                score_val = "no"
                evidence_text = ""
                file_ref = ""

            covered = score_val in ("yes", "partial")
            confidence = 1.0 if score_val == "yes" else 0.5 if score_val == "partial" else 0.0
            evidence = []
            if evidence_text:
                evidence.append(evidence_text)
            if file_ref:
                evidence.append(f"File: {file_ref}")

            criteria_results.append(CriterionResult(
                criterion_id=criterion.id,
                criterion_name=criterion.name,
                covered=covered,
                confidence=confidence,
                evidence=evidence,
            ))

        layer_results.append(LayerResult(
            id=layer.id,
            name=layer.name,
            criteria_results=criteria_results,
        ))

    # Build contract results
    contract_results = []
    for contract in ALL_CONTRACTS:
        criteria_results = []
        for criterion in contract.criteria:
            score_data = scores.get(criterion.id, {})
            if isinstance(score_data, dict):
                score_val = score_data.get("score", "no")
                evidence_text = score_data.get("evidence", "")
                file_ref = score_data.get("file", "")
            elif isinstance(score_data, str):
                score_val = score_data
                evidence_text = ""
                file_ref = ""
            else:
                score_val = "no"
                evidence_text = ""
                file_ref = ""

            covered = score_val in ("yes", "partial")
            confidence = 1.0 if score_val == "yes" else 0.5 if score_val == "partial" else 0.0
            evidence = []
            if evidence_text:
                evidence.append(evidence_text)
            if file_ref:
                evidence.append(f"File: {file_ref}")

            criteria_results.append(CriterionResult(
                criterion_id=criterion.id,
                criterion_name=criterion.name,
                covered=covered,
                confidence=confidence,
                evidence=evidence,
            ))

        contract_results.append(LayerResult(
            id=contract.id,
            name=contract.name,
            criteria_results=criteria_results,
        ))

    return ScanResult(
        project_path=project_path,
        frameworks_detected=fw_dict,
        layer_results=layer_results,
        contract_results=contract_results,
    )


def _empty_result(project_path: str, error: str) -> ScanResult:
    """Return an empty result with an error message."""
    print(f"  Warning: {error}")
    layer_results = []
    for layer in ALL_LAYERS:
        criteria_results = [
            CriterionResult(c.id, c.name, False, 0.0) for c in layer.criteria
        ]
        layer_results.append(LayerResult(layer.id, layer.name, criteria_results))

    contract_results = []
    for contract in ALL_CONTRACTS:
        criteria_results = [
            CriterionResult(c.id, c.name, False, 0.0) for c in contract.criteria
        ]
        contract_results.append(LayerResult(contract.id, contract.name, criteria_results))

    return ScanResult(
        project_path=project_path,
        frameworks_detected={},
        layer_results=layer_results,
        contract_results=contract_results,
    )
