from __future__ import annotations

import json
import os
import re
from pathlib import Path
from typing import Any

import yaml

from .validation import load_schema, validate_context_data


def init_security_context(
    threat_model_path: str | Path,
    output_path: str | Path = "security-context.yaml",
    project_name: str | None = None,
    use_llm: bool = True,
    llm_model: str = "claude-3-5-sonnet-latest",
    schema_path: str | Path | None = None,
) -> dict[str, Any]:
    threat_path = Path(threat_model_path)
    if not threat_path.exists():
        return {
            "ok": False,
            "used_llm": False,
            "output_path": str(output_path),
            "warnings": [],
            "errors": [f"Threat model file does not exist: {threat_path}"],
        }

    threat_model_text = threat_path.read_text(encoding="utf-8")
    warnings: list[str] = []
    used_llm = False

    context: dict[str, Any]
    if use_llm:
        generated = _try_generate_with_llm(
            threat_model_text=threat_model_text,
            project_name=project_name,
            llm_model=llm_model,
        )
        if generated["ok"]:
            context = generated["context"]
            warnings.extend(generated["warnings"])
            used_llm = True
        else:
            warnings.extend(generated["warnings"])
            warnings.append("Falling back to heuristic generator.")
            context = _heuristic_context(threat_model_text, project_name)
    else:
        context = _heuristic_context(threat_model_text, project_name)

    schema = load_schema(schema_path)
    report = validate_context_data(context=context, schema=schema)
    if not report.valid:
        return {
            "ok": False,
            "used_llm": used_llm,
            "output_path": str(output_path),
            "warnings": warnings,
            "errors": [
                f"{err['path']} (line {err['line']}): {err['message']}" for err in report.errors
            ],
        }

    destination = Path(output_path)
    destination.write_text(
        yaml.safe_dump(context, sort_keys=False, allow_unicode=False),
        encoding="utf-8",
    )

    return {
        "ok": True,
        "used_llm": used_llm,
        "output_path": str(destination),
        "warnings": warnings,
        "errors": [],
    }


def _try_generate_with_llm(
    threat_model_text: str,
    project_name: str | None,
    llm_model: str,
) -> dict[str, Any]:
    api_key = os.getenv("ANTHROPIC_API_KEY")
    if not api_key:
        return {
            "ok": False,
            "warnings": [
                "ANTHROPIC_API_KEY is not set; aegis_init cannot use Claude in this environment."
            ],
        }

    try:
        from anthropic import Anthropic
    except Exception:
        return {
            "ok": False,
            "warnings": [
                "anthropic package is not installed; install aegis-mcp[init-llm] for LLM-assisted init."
            ],
        }

    template = {
        "version": "0.2",
        "services": [],
        "trust_boundaries": [],
        "data_classifications": [],
        "approved_patterns": [],
        "anti_patterns": [],
        "authorization_policies": [],
        "input_validations": [],
        "security_configs": [],
        "secrets_management": [],
        "pipeline_policies": [],
        "error_handling": [],
    }

    system_prompt = (
        "You convert threat-model prose into valid Aegis security-context YAML. "
        "Return only JSON with keys: context (object) and unmapped (array of strings)."
    )

    user_prompt = (
        f"Project name: {project_name or 'unknown'}\n"
        "Target contract:\n"
        "- version must be '0.2'\n"
        "- include trust_boundaries, data_classifications, approved_patterns, anti_patterns\n"
        "- include authorization_policies, input_validations, security_configs,\n"
        "  secrets_management, pipeline_policies, and error_handling when applicable\n"
        "- all ids kebab-case\n"
        "- requirements need id/text/severity with severity in low|medium|high|critical\n"
        f"Template:\n{json.dumps(template)}\n\n"
        f"Threat model:\n{threat_model_text}\n"
    )

    client = Anthropic(api_key=api_key)
    message = client.messages.create(
        model=llm_model,
        max_tokens=4096,
        temperature=0,
        system=system_prompt,
        messages=[{"role": "user", "content": user_prompt}],
    )

    response_text = "\n".join(
        part.text for part in message.content if getattr(part, "type", None) == "text"
    ).strip()

    payload = _extract_json_object(response_text)
    context = payload.get("context")
    unmapped = payload.get("unmapped")

    if not isinstance(context, dict):
        return {
            "ok": False,
            "warnings": ["LLM response did not include a valid context object."],
        }

    warnings: list[str] = []
    if isinstance(unmapped, list):
        warnings.extend(str(item) for item in unmapped)

    return {
        "ok": True,
        "context": context,
        "warnings": warnings,
    }


def _extract_json_object(text: str) -> dict[str, Any]:
    stripped = text.strip()
    if stripped.startswith("```"):
        stripped = stripped.strip("`")
        stripped = stripped.replace("json", "", 1).strip()

    try:
        return json.loads(stripped)
    except json.JSONDecodeError:
        pass

    start = stripped.find("{")
    end = stripped.rfind("}")
    if start >= 0 and end > start:
        return json.loads(stripped[start : end + 1])

    raise ValueError("No JSON object found in LLM response")


def _heuristic_context(threat_model_text: str, project_name: str | None) -> dict[str, Any]:
    text = threat_model_text.lower()

    has_public = any(term in text for term in ["public", "internet", "external api", "edge"])
    has_database = any(term in text for term in ["database", "postgres", "mysql", "dynamodb", "sql"])
    has_pii = any(term in text for term in ["pii", "personal data", "email", "address", "dob"])
    has_tokens = any(term in text for term in ["token", "jwt", "oauth", "session"])
    has_auth_logic = any(
        term in text
        for term in ["rbac", "role", "permission", "admin", "owner", "acl", "authorization"]
    )
    has_input_validation = any(
        term in text
        for term in ["upload", "form", "user input", "query parameter", "request body", "url parameter"]
    )
    has_security_config = any(
        term in text for term in ["cors", "header", "cookie", "csp", "hsts", "tls"]
    )
    has_secrets = any(
        term in text
        for term in ["secret", "api key", "credential", "password", "env var", "environment variable"]
    )
    has_pipeline = bool(
        re.search(r"\b(ci|cd|pipeline|deploy|jenkins|buildkite)\b", text)
        or "github actions" in text
        or "gitlab ci" in text
    )
    has_error_handling = any(
        term in text for term in ["error", "exception", "stack trace", "500", "error handling"]
    )

    services = _extract_services(text)

    trust_boundaries: list[dict[str, Any]] = []
    if has_public:
        trust_boundaries.append(
            {
                "id": "public-to-internal",
                "name": "Public traffic to internal services",
                "from_components": ["internet"],
                "to_components": ["application"],
                "crossing_requirements": [
                    {
                        "id": "enforce-authn",
                        "text": "Public-facing requests must enforce authentication and authorization checks.",
                        "severity": "critical",
                    },
                    {
                        "id": "rate-limit-public",
                        "text": "Public-facing endpoints must enforce rate limits and abuse controls.",
                        "severity": "high",
                    },
                ],
            }
        )

    if has_database:
        trust_boundaries.append(
            {
                "id": "internal-to-database",
                "name": "Internal services to data stores",
                "from_components": ["application"],
                "to_components": ["database"],
                "crossing_requirements": [
                    {
                        "id": "least-privileged-db-role",
                        "text": "Database access must use least-privileged application roles.",
                        "severity": "high",
                    },
                    {
                        "id": "tls-db-transport",
                        "text": "Database connections must require encrypted transport.",
                        "severity": "high",
                    },
                ],
            }
        )

    if not trust_boundaries:
        trust_boundaries.append(
            {
                "id": "internal-only",
                "name": "Internal processing boundary",
                "from_components": ["internal-service"],
                "to_components": ["internal-service"],
                "crossing_requirements": [
                    {
                        "id": "secure-defaults",
                        "text": "Internal systems must use secure defaults and principle-of-least-privilege.",
                        "severity": "medium",
                    }
                ],
            }
        )

    data_classifications: list[dict[str, Any]] = []

    if has_pii:
        data_classifications.append(
            {
                "id": "user-pii",
                "name": "User PII",
                "description": "Personally identifiable data related to users.",
                "examples": ["email", "name", "address"],
                "handling_rules": [
                    {
                        "id": "no-pii-logs",
                        "text": "PII must be redacted before logging.",
                        "severity": "critical",
                    },
                    {
                        "id": "encrypt-pii-at-rest",
                        "text": "PII must be encrypted at rest.",
                        "severity": "high",
                    },
                ],
            }
        )

    if has_tokens:
        data_classifications.append(
            {
                "id": "auth-token",
                "name": "Authentication tokens",
                "description": "Bearer tokens or session artifacts used for auth.",
                "examples": ["jwt", "refresh-token"],
                "handling_rules": [
                    {
                        "id": "no-token-logs",
                        "text": "Authentication tokens must never be logged in plaintext.",
                        "severity": "critical",
                    }
                ],
            }
        )

    if not data_classifications:
        data_classifications.append(
            {
                "id": "internal-data",
                "name": "Internal data",
                "description": "General internal system data.",
                "examples": [],
                "handling_rules": [
                    {
                        "id": "minimum-retention",
                        "text": "Retain only the minimum data required for business operation.",
                        "severity": "medium",
                    }
                ],
            }
        )

    authorization_policies: list[dict[str, Any]] = []
    if has_auth_logic or has_public:
        authorization_policies.append(
            {
                "id": "resource-owner",
                "name": "Resource Owner Authorization",
                "description": "Mutating and read-sensitive operations must verify ownership or explicit grant.",
                "required_controls": [
                    {
                        "id": "enforce-ownership-check",
                        "text": "Resource-level access must verify caller ownership or an equivalent scoped grant.",
                        "severity": "critical",
                    },
                    {
                        "id": "deny-by-default-authz",
                        "text": "Authorization decisions must fail closed when policy metadata is unavailable.",
                        "severity": "high",
                    },
                ],
            }
        )
        authorization_policies.append(
            {
                "id": "admin-only",
                "name": "Administrative Access Control",
                "description": "Administrative capabilities require explicit privileged role checks.",
                "required_controls": [
                    {
                        "id": "enforce-admin-role-check",
                        "text": "Administrative operations must enforce an explicit privileged role check on the server.",
                        "severity": "critical",
                    }
                ],
            }
        )

    input_validations: list[dict[str, Any]] = []
    if has_input_validation or has_public:
        input_validations.append(
            {
                "id": "request-body",
                "name": "Request Body Validation",
                "description": "Untrusted body payloads must be schema validated before use.",
                "validation_rules": [
                    {
                        "id": "validate-request-schema",
                        "text": "Request bodies must be validated against explicit schemas with strict type checks.",
                        "severity": "high",
                    },
                    {
                        "id": "reject-unknown-fields",
                        "text": "Unknown or unexpected body fields must be rejected by default.",
                        "severity": "medium",
                    },
                ],
            }
        )
        input_validations.append(
            {
                "id": "url-parameter",
                "name": "URL and Query Parameter Validation",
                "description": "Route and query parameters must be constrained and canonicalized.",
                "validation_rules": [
                    {
                        "id": "validate-url-parameters",
                        "text": "Route and query parameters must be validated for type, format, and length before use.",
                        "severity": "high",
                    }
                ],
            }
        )

    if "upload" in text or "file upload" in text:
        input_validations.append(
            {
                "id": "file-upload",
                "name": "File Upload Validation",
                "description": "Uploaded files must be validated and constrained before processing.",
                "validation_rules": [
                    {
                        "id": "allowlist-upload-types",
                        "text": "File uploads must enforce MIME type and extension allowlists before storage.",
                        "severity": "high",
                    },
                    {
                        "id": "enforce-upload-size-limits",
                        "text": "File uploads must enforce strict size limits and reject oversized payloads.",
                        "severity": "medium",
                    },
                ],
            }
        )

    security_configs: list[dict[str, Any]] = []
    if has_security_config or has_public:
        security_configs.append(
            {
                "id": "cors-policy",
                "name": "CORS Policy",
                "description": "Cross-origin access must be origin-allowlisted with credential safety controls.",
                "scope": "edge",
                "when": {
                    "boundary_crossings": ["public-to-internal"],
                    "actions": ["create-endpoint", "external-api-call"],
                },
                "rules": [
                    {
                        "id": "restrict-cors-origins",
                        "text": "CORS must use explicit origin allowlists; wildcard origins are prohibited for credentialed requests.",
                        "severity": "high",
                    }
                ],
            }
        )
        security_configs.append(
            {
                "id": "security-headers",
                "name": "HTTP Security Headers",
                "description": "Public HTTP responses must include hardened baseline security headers.",
                "scope": "edge",
                "when": {
                    "boundary_crossings": ["public-to-internal"],
                    "actions": ["create-endpoint"],
                },
                "rules": [
                    {
                        "id": "set-security-headers",
                        "text": "Public responses must set HSTS, X-Content-Type-Options, and frame-ancestors controls where applicable.",
                        "severity": "medium",
                    }
                ],
            }
        )

    if has_tokens:
        security_configs.append(
            {
                "id": "cookie-settings",
                "name": "Session Cookie Settings",
                "description": "Session cookies must be configured with secure transport and CSRF-resistant defaults.",
                "scope": "application",
                "when": {"actions": ["auth-flow-change", "create-endpoint"]},
                "rules": [
                    {
                        "id": "secure-session-cookies",
                        "text": "Session cookies must be Secure, HttpOnly, and SameSite=strict or equivalent mitigation.",
                        "severity": "high",
                    }
                ],
            }
        )

    secrets_management: list[dict[str, Any]] = []
    if has_secrets or has_tokens:
        secrets_management.append(
            {
                "id": "no-hardcoded-secrets",
                "name": "No Hardcoded Secrets",
                "description": "Secrets must never be embedded in source, tests, or configuration defaults.",
                "when": {},
                "rules": [
                    {
                        "id": "store-secrets-in-manager",
                        "text": "Secrets must be stored in approved secret managers or protected runtime environment variables.",
                        "severity": "critical",
                    }
                ],
            }
        )
        secrets_management.append(
            {
                "id": "secret-rotation",
                "name": "Secret Rotation",
                "description": "Long-lived credentials must be rotated and revocable.",
                "when": {},
                "rules": [
                    {
                        "id": "rotate-secrets-regularly",
                        "text": "Production secrets and API keys must have automated or scheduled rotation with revocation support.",
                        "severity": "high",
                    }
                ],
            }
        )

    pipeline_policies: list[dict[str, Any]] = []
    if has_pipeline:
        pipeline_policies.append(
            {
                "id": "dependency-scanning",
                "name": "Dependency Scanning in CI",
                "description": "Build pipelines must block merges for known exploitable dependency vulnerabilities.",
                "when": {},
                "rules": [
                    {
                        "id": "run-dependency-vulnerability-scan",
                        "text": "CI must run dependency vulnerability scanning on pull requests and block on high/critical findings.",
                        "severity": "high",
                    }
                ],
            }
        )
        pipeline_policies.append(
            {
                "id": "signed-commits",
                "name": "Signed Commits for Protected Branches",
                "description": "Protected branches require signed commits and verified identities.",
                "when": {},
                "rules": [
                    {
                        "id": "require-signed-commits",
                        "text": "Protected branches must enforce commit signature verification before merge.",
                        "severity": "medium",
                    }
                ],
            }
        )

    error_handling: list[dict[str, Any]] = []
    if has_error_handling or has_public:
        error_handling.append(
            {
                "id": "safe-error-responses",
                "name": "Safe Error Responses",
                "description": "Externally visible errors must avoid sensitive implementation details.",
                "when": {"actions": ["create-endpoint", "auth-flow-change"]},
                "rules": [
                    {
                        "id": "no-sensitive-error-details",
                        "text": "Error responses must not disclose secrets, SQL fragments, internal hostnames, or stack details.",
                        "severity": "high",
                    }
                ],
            }
        )
        error_handling.append(
            {
                "id": "no-stack-trace-leak",
                "name": "No Stack Trace Leaks",
                "description": "Unhandled exceptions must not expose stack traces to clients.",
                "when": {"actions": ["create-endpoint"]},
                "rules": [
                    {
                        "id": "suppress-stack-trace-responses",
                        "text": "Unhandled exceptions must return sanitized errors and log full traces only in protected sinks.",
                        "severity": "high",
                    }
                ],
            }
        )

    context = {
        "version": "0.2",
        "services": services,
        "trust_boundaries": trust_boundaries,
        "data_classifications": data_classifications,
        "approved_patterns": [
            {
                "id": "auth-middleware",
                "name": "Authentication Middleware",
                "description": "Use shared authentication middleware for request validation.",
                "example": "router.use(authenticate())",
                "when": {
                    "boundary_crossings": ["public-to-internal"],
                    "actions": ["create-endpoint", "auth-flow-change"],
                },
                "requirements": [],
            },
            {
                "id": "parameterized-queries",
                "name": "Parameterized Queries",
                "description": "Use bound parameters for all database queries.",
                "example": "db.execute('SELECT * FROM t WHERE id = ?', [id])",
                "when": {
                    "actions": ["database-query"],
                },
                "requirements": [],
            },
        ],
        "anti_patterns": [
            {
                "id": "raw-sql",
                "name": "Raw SQL Interpolation",
                "description": "Do not concatenate untrusted input into SQL statements.",
                "severity": "critical",
                "when": {"actions": ["database-query"]},
                "requirements": [
                    {
                        "id": "no-sql-interpolation",
                        "text": "SQL statements must use parameter binding instead of interpolation.",
                        "severity": "critical",
                    }
                ],
            },
            {
                "id": "secrets-in-logs",
                "name": "Secrets in Logs",
                "description": "Never emit credentials, tokens, or secrets to logs.",
                "severity": "critical",
                "when": {},
                "requirements": [],
            },
        ],
        "authorization_policies": authorization_policies,
        "input_validations": input_validations,
        "security_configs": security_configs,
        "secrets_management": secrets_management,
        "pipeline_policies": pipeline_policies,
        "error_handling": error_handling,
    }

    return context


def _extract_services(text: str) -> list[dict[str, str]]:
    candidates = set()

    for match in re.finditer(r"\b([a-z0-9][a-z0-9-]{2,30})(?:\s+service|\s+api)\b", text):
        service_id = _kebab(match.group(1))
        if service_id:
            candidates.add(service_id)

    for match in re.finditer(r"\b([a-z0-9][a-z0-9-_]{2,30})\b", text):
        token = match.group(1)
        if token.endswith(("-api", "-service")):
            candidates.add(_kebab(token))

    services = []
    for service_id in sorted(candidates):
        if service_id in {"public", "internal", "service", "api"}:
            continue
        services.append(
            {
                "id": service_id,
                "name": service_id.replace("-", " ").title(),
                "description": "Auto-detected from threat model text.",
            }
        )

    return services


def _kebab(value: str) -> str:
    normalized = re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-")
    normalized = re.sub(r"-+", "-", normalized)
    return normalized
