import { parse } from "@babel/parser";
import _traverse from "@babel/traverse";
import * as t from "@babel/types";
import {
  makeStyleDeclarationFact,
  makeStyleUnsupportedFact,
  type Fact,
  type FactLocation,
  type StyleDeclarationFact,
  type UnsupportedStyleReason,
} from "@fragments-sdk/core";

import { BABEL_PARSER_OPTIONS } from "./babel-config.js";
import { extractScssFacts } from "./scss.js";

const traverse = (_traverse as unknown as { default?: typeof _traverse }).default ?? _traverse;

interface CssInJsTag {
  label: string;
  syntax: string;
  location: FactLocation;
}

interface ImportNames {
  styled: Set<string>;
  css: Set<string>;
  style: Set<string>;
}

export function extractCssInJsFacts(source: string, filePath: string): Fact[] {
  let ast: t.File;
  try {
    ast = parse(source, BABEL_PARSER_OPTIONS);
  } catch {
    return [];
  }

  const imports = collectCssInJsImports(ast);
  const constLiterals = collectStaticConstLiterals(ast);

  const facts: Fact[] = [];
  traverse(ast, {
    TaggedTemplateExpression(path) {
      const tag = cssInJsTag(path.node.tag, imports, filePath);
      if (!tag) return;
      const { css, unsupported } = staticTemplateCss(
        path.node.quasi,
        filePath,
        tag,
        constLiterals,
      );
      if (!css.trim()) return;
      facts.push(
        ...cssFactsFromTemplate(
          css,
          filePath,
          tag,
          locFromNode(path.node.quasi.quasis[0], filePath) ?? tag.location,
        ),
      );
      facts.push(...unsupported);
    },
    CallExpression(path) {
      const firstArg = path.node.arguments[0];
      const tag = cssInJsObjectCallTag(path.node.callee, imports, filePath);
      if (!tag) return;
      if (t.isObjectExpression(firstArg)) {
        facts.push(...cssFactsFromObject(firstArg, filePath, tag));
      } else if (t.isArrayExpression(firstArg)) {
        facts.push(...cssFactsFromArray(firstArg, filePath, tag));
      }
    },
    JSXAttribute(path) {
      if (
        !t.isJSXIdentifier(path.node.name) ||
        (path.node.name.name !== "css" && path.node.name.name !== "sx")
      ) {
        return;
      }
      const value = path.node.value;
      if (!t.isJSXExpressionContainer(value)) {
        return;
      }
      const location = locFromNode(path.node.name, filePath);
      if (!location) return;
      const tag = {
        label: `${path.node.name.name}-prop`,
        syntax: `${path.node.name.name}-prop`,
        location,
      };
      if (t.isObjectExpression(value.expression)) {
        facts.push(...cssFactsFromObject(value.expression, filePath, tag));
      } else if (t.isArrayExpression(value.expression)) {
        facts.push(...cssFactsFromArray(value.expression, filePath, tag));
      } else {
        facts.push(
          unsupportedStyleFact({
            filePath,
            tag,
            declarationPath: "jsx-attribute",
            reason: "dynamic-object-value",
            location,
          }),
        );
      }
    },
  });
  return facts;
}

function collectCssInJsImports(ast: t.File): ImportNames {
  const names: ImportNames = {
    styled: new Set(),
    css: new Set(),
    style: new Set(),
  };
  for (const node of ast.program.body) {
    if (!t.isImportDeclaration(node)) continue;
    const source = node.source.value;
    for (const specifier of node.specifiers) {
      if (
        (source === "styled-components" || source === "@emotion/styled") &&
        t.isImportDefaultSpecifier(specifier)
      ) {
        names.styled.add(specifier.local.name);
      }
      if (
        source === "@stitches/react" &&
        t.isImportSpecifier(specifier) &&
        t.isIdentifier(specifier.imported) &&
        specifier.imported.name === "styled"
      ) {
        names.styled.add(specifier.local.name);
      }
      if (
        (source === "@emotion/react" || source === "@stitches/react") &&
        t.isImportSpecifier(specifier) &&
        t.isIdentifier(specifier.imported) &&
        specifier.imported.name === "css"
      ) {
        names.css.add(specifier.local.name);
      }
      if (
        source === "@vanilla-extract/css" &&
        t.isImportSpecifier(specifier) &&
        t.isIdentifier(specifier.imported) &&
        specifier.imported.name === "style"
      ) {
        names.style.add(specifier.local.name);
      }
    }
  }
  return names;
}

function collectStaticConstLiterals(ast: t.File): Map<string, string> {
  const out = new Map<string, string>();
  for (const node of ast.program.body) {
    if (!t.isVariableDeclaration(node) || node.kind !== "const") continue;
    for (const declarator of node.declarations) {
      if (!t.isIdentifier(declarator.id) || !declarator.init) continue;
      const value = staticLiteralValue(declarator.init);
      if (value !== null) out.set(declarator.id.name, value);
    }
  }
  return out;
}

function cssInJsTag(
  tag: t.Expression,
  imports: ImportNames,
  filePath: string,
): CssInJsTag | null {
  const location = locFromNode(tag, filePath);
  if (!location) return null;

  if (t.isIdentifier(tag) && imports.css.has(tag.name)) {
    return { label: tag.name, syntax: "css-tag", location };
  }
  const styled = styledExpressionLabel(tag, imports);
  if (styled) {
    return { label: styled, syntax: "styled-tag", location };
  }
  return null;
}

function cssInJsObjectCallTag(
  callee: t.Expression | t.V8IntrinsicIdentifier,
  imports: ImportNames,
  filePath: string,
): CssInJsTag | null {
  if (t.isV8IntrinsicIdentifier(callee)) return null;
  const location = locFromNode(callee, filePath);
  if (!location) return null;

  if (t.isIdentifier(callee) && imports.css.has(callee.name)) {
    return { label: callee.name, syntax: "css-object", location };
  }
  if (t.isIdentifier(callee) && imports.style.has(callee.name)) {
    return { label: callee.name, syntax: "vanilla-extract-style", location };
  }
  const styled = styledExpressionLabel(callee, imports);
  if (styled && /\.(?:attrs|withConfig)$/.test(styled)) return null;
  if (styled) {
    return { label: styled, syntax: "styled-object", location };
  }
  return null;
}

function styledExpressionLabel(
  expression: t.Expression | t.V8IntrinsicIdentifier,
  imports: ImportNames,
): string | null {
  if (t.isV8IntrinsicIdentifier(expression)) return null;
  if (t.isIdentifier(expression) && imports.styled.has(expression.name)) {
    return expression.name;
  }
  if (t.isMemberExpression(expression)) {
    const object = styledExpressionLabel(expression.object, imports);
    if (!object) return null;
    return `${object}.${memberName(expression.property) ?? "unknown"}`;
  }
  if (t.isCallExpression(expression)) {
    const callee = styledExpressionLabel(expression.callee, imports);
    return callee ? `${callee}()` : null;
  }
  return null;
}

function staticTemplateCss(
  template: t.TemplateLiteral,
  filePath: string,
  tag: CssInJsTag,
  constLiterals: ReadonlyMap<string, string>,
): { css: string; unsupported: Fact[] } {
  const unsupported: Fact[] = [];
  const css = template.quasis
    .map((quasi, index) => {
      const expression = template.expressions[index];
      if (!expression) return quasi.value.raw;
      const value = staticExpressionValue(expression, constLiterals);
      if (value !== null) return `${quasi.value.raw}${value}`;
      unsupported.push(
        unsupportedStyleFact({
          filePath,
          tag,
          declarationPath: `template/${index}`,
          reason: "dynamic-template-expression",
          location: locFromNode(expression, filePath) ?? tag.location,
        }),
      );
      return `${quasi.value.raw}/*__expr__*/`;
    })
    .join("");
  return { css, unsupported };
}

function cssFactsFromTemplate(
  css: string,
  filePath: string,
  tag: CssInJsTag,
  templateLocation: FactLocation,
): Fact[] {
  return extractScssFacts(css, filePath)
    .filter((fact): fact is StyleDeclarationFact => fact.kind === "style_declaration")
    .map((fact) =>
      makeStyleDeclarationFact({
        file: filePath,
        selector: `css-in-js:${tag.label}:${tag.location.line}:${fact.selector}`,
        declarationPath: fact.declarationPath,
        property: fact.property,
        value: fact.value,
        location: offsetTemplateLocation(fact.location, templateLocation),
      }),
    );
}

function cssFactsFromObject(
  object: t.ObjectExpression,
  filePath: string,
  tag: CssInJsTag,
  selector = "",
  prefix = "",
): Fact[] {
  const facts: Fact[] = [];
  object.properties.forEach((property, index) => {
    const declarationPath = prefix ? `${prefix}/${index}` : String(index);
    if (t.isSpreadElement(property)) {
      facts.push(
        unsupportedStyleFact({
          filePath,
          tag,
          selector,
          declarationPath,
          reason: "object-spread",
          location: locFromNode(property, filePath) ?? tag.location,
        }),
      );
      return;
    }
    if (!t.isObjectProperty(property) || property.computed) return;
    const key = objectKey(property.key);
    if (!key) return;
    if (t.isObjectExpression(property.value)) {
      facts.push(
        ...cssFactsFromObject(
          property.value,
          filePath,
          tag,
          selector ? `${selector} ${key}` : key,
          declarationPath,
        ),
      );
      return;
    }

    const value = staticObjectStyleValue(property.value);
    const location = locFromNode(property.key, filePath);
    if (!location) return;
    if (value === null) {
      facts.push(
        unsupportedStyleFact({
          filePath,
          tag,
          selector,
          declarationPath,
          property: cssPropertyName(key),
          reason: "dynamic-object-value",
          location,
        }),
      );
      return;
    }
    facts.push(
      makeStyleDeclarationFact({
        file: filePath,
        selector: `css-in-js:${tag.label}:${tag.location.line}:${selector}`,
        declarationPath,
        property: cssPropertyName(key),
        value,
        location,
      }),
    );
  });
  return facts;
}

function cssFactsFromArray(
  array: t.ArrayExpression,
  filePath: string,
  tag: CssInJsTag,
  selector = "",
  prefix = "array",
): Fact[] {
  const facts: Fact[] = [];
  array.elements.forEach((element, index) => {
    const declarationPath = `${prefix}/${index}`;
    if (t.isObjectExpression(element)) {
      facts.push(
        ...cssFactsFromObject(
          element,
          filePath,
          tag,
          selector,
          declarationPath,
        ),
      );
      return;
    }
    if (element === null) return;
    facts.push(
      unsupportedStyleFact({
        filePath,
        tag,
        selector,
        declarationPath,
        reason: "unsupported-array-element",
        location: locFromNode(element, filePath) ?? tag.location,
      }),
    );
  });
  return facts;
}

function objectKey(key: t.ObjectProperty["key"]): string | null {
  if (t.isIdentifier(key)) return key.name;
  if (t.isStringLiteral(key)) return key.value;
  if (t.isNumericLiteral(key)) return String(key.value);
  return null;
}

function staticObjectStyleValue(value: t.ObjectProperty["value"]): string | null {
  return staticLiteralValue(value);
}

function staticExpressionValue(
  value: t.Expression | t.TSType,
  constLiterals: ReadonlyMap<string, string>,
): string | null {
  if (t.isIdentifier(value)) return constLiterals.get(value.name) ?? null;
  return staticLiteralValue(value);
}

function staticLiteralValue(value: t.Node): string | null {
  if (t.isStringLiteral(value)) return value.value;
  if (t.isNumericLiteral(value)) return String(value.value);
  if (t.isTemplateLiteral(value) && value.expressions.length === 0) {
    return value.quasis.map((quasi) => quasi.value.raw).join("");
  }
  if (
    t.isUnaryExpression(value) &&
    (value.operator === "-" || value.operator === "+") &&
    t.isNumericLiteral(value.argument)
  ) {
    return `${value.operator}${value.argument.value}`;
  }
  return null;
}

function unsupportedStyleFact(args: {
  filePath: string;
  tag: CssInJsTag;
  selector?: string;
  declarationPath: string;
  property?: string;
  reason: UnsupportedStyleReason;
  location: FactLocation;
}): Fact {
  return makeStyleUnsupportedFact({
    file: args.filePath,
    source: "css-in-js",
    syntax: args.tag.syntax,
    selector: `css-in-js:${args.tag.label}:${args.tag.location.line}:${
      args.selector ?? ""
    }`,
    declarationPath: args.declarationPath,
    property: args.property,
    reason: args.reason,
    location: args.location,
  });
}

function cssPropertyName(value: string): string {
  if (value.startsWith("--") || value.includes("-")) return value;
  return value.replace(/[A-Z]/g, (char) => `-${char.toLowerCase()}`);
}

function offsetTemplateLocation(
  location: FactLocation,
  tagLocation: FactLocation,
): FactLocation {
  return {
    file: location.file,
    line: tagLocation.line + location.line - 1,
    column:
      location.line === 1
        ? tagLocation.column + location.column
        : location.column,
  };
}

function locFromNode(node: t.Node, file: string): FactLocation | null {
  if (!node.loc) return null;
  return {
    file,
    line: node.loc.start.line,
    column: node.loc.start.column,
  };
}

function memberName(node: t.Expression | t.PrivateName): string | null {
  if (t.isIdentifier(node)) return node.name;
  if (t.isStringLiteral(node)) return node.value;
  return null;
}
